接上一篇安装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
调用后:
参考: