pytorch设置gpu进行训练

接上一篇安装gpu版本pytorch后，这篇描述设置gpu进行训练

（1）模型设置

cuda_gpu = torch.cuda.is_available()
if(cuda_gpu):
    net = torch.nn.DataParallel(net).cuda()

 

（2）输入数据设置

.cuda()将tensor专程cuda

inputs, labels = Variable(inputs.cuda()), Variable(labels.cuda())

 

（3）输出设置

.cpu将cuda转成tensor

outputs = net(Variable(images.cuda()))
_, predicted = torch.max(outputs.data, 1)
countall = countall + len(labels.cuda())
for j in range(len(labels)):
    if(classes[labels.cuda()[j]] == classes[predicted.cuda()[j]]):
        countok = countok + 1
print("countall:",countall,", countok:",countok,",Acc:",countok/countall)

 

最后，查看nvidia-smi中，gpu被调用。

未调用前，只有这个gpu pid

调用后：

参考：

https://www.cnblogs.com/darkknightzh/p/6836568.html

https://blog.csdn.net/qq_21578849/article/details/85240797