GPU 加速运算
作者:
编辑:
莫烦
2017-05-27
学习资料:
要点 ¶
在 GPU 训练可以大幅提升运算速度. 而且 Torch 也有一套很好的 GPU 运算体系. 但是要强调的是:
- 你的电脑里有合适的 GPU 显卡(NVIDIA), 且支持 CUDA 模块. 请在NVIDIA官网查询
- 必须安装 GPU 版的 Torch, 点击这里查看如何安装
用 GPU 训练 CNN ¶
这份 GPU 的代码是依据之前这份CNN的代码修改的.
大概修改的地方包括将数据的形式变成 GPU 能读的形式, 然后将 CNN 也变成 GPU 能读的形式. 做法就是在后面加上 .cuda(), 很简单.
...
test_data = torchvision.datasets.MNIST(root='./mnist/', train=False)
# !!!!!!!! 修改 test data 形式 !!!!!!!!! #
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(torch.FloatTensor)[:2000].cuda()/255. # Tensor on GPU
test_y = test_data.test_labels[:2000].cuda()
再来把我们的 CNN 参数也变成 GPU 兼容形式.
class CNN(nn.Module):
...
cnn = CNN()
# !!!!!!!! 转换 cnn 去 CUDA !!!!!!!!! #
cnn.cuda() # Moves all model parameters and buffers to the GPU.
然后就是在 train 的时候, 将每次的training data 变成 GPU 形式. + .cuda()
for epoch ..:
for step, ...:
# !!!!!!!! 这里有修改 !!!!!!!!! #
b_x = x.cuda() # Tensor on GPU
b_y = y.cuda() # Tensor on GPU
...
if step % 50 == 0:
test_output = cnn(test_x)
# !!!!!!!! 这里有修改 !!!!!!!!! #
pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze() # 将操作放去 GPU
accuracy = torch.sum(pred_y == test_y) / test_y.size(0)
...
test_output = cnn(test_x[:10])
# !!!!!!!! 这里有修改 !!!!!!!!! #
pred_y = torch.max(test_output, 1)[1].cuda().data.squeeze() # 将操作放去 GPU
...
print(test_y[:10], 'real number')
大功告成~
所以这也就是在我 github 代码 中的每一步的意义啦.
转移至 CPU ¶
如果你有些计算还是需要在 CPU 上进行的话呢, 比如 plt 的可视化, 我们需要将这些计算或者数据转移至 CPU.
cpu_data = gpu_data.cpu()
分享到:
如果你觉得这篇文章或视频对你的学习很有帮助, 请你也分享它, 让它能再次帮助到更多的需要学习的人.
莫烦没有正式的经济来源, 如果你也想支持 莫烦Python 并看到更好的教学内容, 赞助他一点点, 作为鼓励他继续开源的动力.
