切换视频源:

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

降低知识传递的门槛

莫烦的对教育的态度是: 降低知识传递的门槛,不希望给"学习"设置金钱障碍。 这是我花大量业余时间贡献 免费 AI分享的原因。 通过 【赞助】 能及时让我看到你对 【莫烦态度】 的认同,我也更有理由坚持下去。

如果你当前目标是找工作或者转行AI,想接受更加丰富的教学资源、培训辅导体验,我想推荐我的朋友 七月在线 给你, 通过这个 【莫烦Python为你提供的注册链接】, 你将可以获得莫烦专门为你协商的课程优惠券。祝你找/换工作顺利~