模型放到gpu上训练
模型放到gpu上训练
·
把模型放到gpu上训练只需要设置三个地方即可:
(1) 指定训练设备device;
device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
torch.device()
就是装torch.Tensor
的一个地方 ;
'cuda'
这里如果没有指定具体的卡号,那么系统默认 cuda:0;
(2)实例化模型时候把模型加到device;
model = NN(input_size, num_classes).to(device)
(3) 训练时候把特征和标签加到device, 测试时候也同样需要加到device;
data = data.to(device)
targets = targets.to(device)
在多卡上并行计算
torch.nn.DataParallel()
torch.nn.DataParallel()
具体的过程:大体就是将模型加载的每个卡上,数据平均分到每个卡上,原则上保证batch_size大于卡的数目就行;
device = torch.device('cuda:2') #device = torch.device("cuda:1" if use_cuda else "cpu")
model = resnet19()
if torch.cuda.device_count() > 1: #10
print(torch.cuda.device_count())
model = nn.DataParallel(model, device_ids = [2,3,4])
model.to(device)
这段代码运行之后占用的GPU是:0,2,3,4。为什么会占用0?
因为 即使我指定的卡没有0卡,他也会在0卡里面放参数和缓存;
如何避免这种现象呢?
改变默认的device_ids[0]
指向的卡。默认device_ids[0]
指向的就是0卡,只需要通过环境变量,让device_ids[0]
指向其他卡片即可;
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3, 4" #仅有2,3,4(物理上)卡对程序可见,对应的逻辑上的卡号就算0,1,2
......
device = torch.device('cuda') #device = torch.device("cuda:1" if use_cuda else "cpu")
model = resnet19()
if torch.cuda.device_count() > 1: #10
print(torch.cuda.device_count())
model = nn.DataParallel(model, device_ids = [0,1,2])
model.to(device)
参考:Pytorch | GPU | 将代码加载到GPU上运行_sunflower_level1的博客-CSDN博客_pytorch在gpu上运行
做记录方便自己看,希望大家去看原作者的,更简单易懂!!!
更多推荐
所有评论(0)