Pytorch：使用GPU训练

本篇文章接着**上篇文章** 的模型继续探讨如何使用GPU训练。使用GPU 训练可以大幅提升运算速度，降低时间成本，提高深度学习的体验感。 Pytorch 有一套很好的支持 GPU 运算体系，因此使用起来非常方便。
在代码中使用GPU训练主要有三处需要注意：模型转为cuda，数据转为cuda，和输出数据去cuda，转为numpy。修改的地方包括将数据的形式变成 GPU 能读的形式, 然后将 网络模型 也变成 GPU 能读的形式。做法就是在后面加上 .cuda() 。

代码传送门

1、模型转为cuda

gpus = [0]   #使用哪几个GPU进行训练，这里选择0号GPU
cuda_gpu = torch.cuda.is_available()   #判断GPU是否存在可用
net = Net(12288, 25, 16, 6)
if(cuda_gpu):
    net = torch.nn.DataParallel(net, device_ids=gpus).cuda()   #将模型转为cuda类型

2、数据转为cuda

(minibatchX, minibatchY) = minibatch
minibatchX = minibatchX.astype(np.float32).T
minibatchY = minibatchY.astype(np.float32).T
if(cuda_gpu):
    b_x = Variable(torch.from_numpy(minibatchX).cuda())    #将数据转为cuda类型
    b_y = Variable(torch.from_numpy(minibatchY).cuda())
else:
    b_x = Variable(torch.from_numpy(minibatchX))
    b_y = Variable(torch.from_numpy(minibatchY))

3、输出数据去cuda，转为numpy

correct_prediction = sum(torch.max(output, 1)[1].data.squeeze() == torch.max(b_y, 1)[1].data.squeeze())
if(cuda_gpu):
    correct_prediction = correct_prediction.cpu().numpy()   #.cpu将cuda转为tensor类型，.numpy将tensor转为numpy类型
else:
    correct_prediction = correct_prediction.numpy()

输入nvidia-smi，可以看到调用GPU成功！