Pytorch Multi GPU 使用

1 检测torch 是否支持gpu

import torch
print(torch.cuda.version) #查看cuda版本
print(torch.cuda.is_available()) #查看pytorch 是否能用 cuda
--9.0.176
--True

如果输出为true 就为可用，如果只有版本号，且输出为False，则需要重新装一下显卡的驱动

2 前向计算

device_ids = [0，1] #这里看你有几个GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = model(input) 
net = nn.DataParallel(net, device_ids=device_ids)
model.to(device) # 将模型放到GPU-0上，可以查看GPU的显存增加了一定的量 ，
##没有这句会报错，
##raise TypeError('Broadcast function not implemented for CPU tensors')

3 后向计算

optimizer = optim.SGD(param, momentum=momentum, weight_decay=weight_decay)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)  # 将optim也并行

4 训练

predicted = net(img) 
loss_label = loss(label, predicted) 
 
# compute gradient and do SGD step
optimizer.zero_grad() 
loss_label .backward() # 如果loss函数是自己写的，输出的维度是GPU个数的loss，所以需要合并成同一个维度，loss.sum(). .backward() 
#直接用 loss_label .backward()  会报错
#grad can be implicitly created only for scalar outputs 
optimizer.module.step()