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()