Question

作者: Morvan 编辑: Morvan

• 学习资料:
  • 本节的全部代码
  • PyTorch 官网

在 GPU 训练可以大幅提升运算速度. 而且 Torch 也有一套很好的 GPU 运算体系. 但是要强调的是:

• 你的电脑里有合适的 GPU 显卡(NVIDIA), 且支持 CUDA 模块。 请在 NVIDIA 官网查询
• 必须安装 GPU 版的 Torch, 点击这里查看如何安装

本节内容包括:

• 用 GPU 训练 CNN
• 动态 RNN

用 GPU 训练 CNN

这份 GPU 的代码是依据 之前这份 CNN 的代码修改的. 大概修改的地方包括将数据的形式变成 GPU 能读的形式，然后将 CNN 也变成 GPU 能读的形式. 做法就是在后面加上 .cuda() , 很简单。

...

test_data = torchvision.datasets.MNIST(root='./mnist/', train=False)

# !!!!!!!! 修改 test data 形式 !!!!!!!!! #
test_x = Variable(torch.unsqueeze(test_data.test_data, dim=1)).type(torch.FloatTensor)[:2000].cuda()/255.   # Tensor on GPU
test_y = test_data.test_labels[:2000]

再来把我们的 CNN 参数也变成 GPU 兼容形式。

class CNN(nn.Module):
    ...

cnn = CNN()

# !!!!!!!! 转换 cnn !!!!!!!!! #
cnn.cuda()      # Moves all model parameters and buffers to the GPU.

然后就是在 train 的时候，将每次的 training data 变成 GPU 形式. + .cuda()

for epoch ..:
    for step, ...:
        # !!!!!!!! 转换批数据 !!!!!!!!! #
        b_x = Variable(x).cuda()    # Tensor on GPU
        b_y = Variable(y).cuda()    # Tensor on GPU

        ...

        if step % 50 == 0:
            test_output = cnn(test_x)

            # !!!!!!!! 将结果放到 CPU 进行逻辑分析 !!!!!!!!! #
            pred_y = torch.max(test_output, 1)[1].cup().data.squeeze()  # Move to CPU

            accuracy = sum(pred_y == test_y) / test_y.size(0)
            ...

test_output = cnn(test_x[:10])

# !!!!!!!! 将结果放到 CPU 进行逻辑分析 !!!!!!!!! #
pred_y = torch.max(test_output, 1)[1].cup().data.numpy().squeeze()  # Move to CPU
...

大功告成~

所以这也就是在我 github 代码 中的每一步的意义啦。

如果你觉得这篇文章或视频对你的学习很有帮助，请你也分享它，让它能再次帮助到更多的需要学习的人。