GPU 加速运算

学习资料:

• 本节的全部代码
• PyTorch 官网

要点¶

在 GPU 训练可以大幅提升运算速度. 而且 Torch 也有一套很好的 GPU 运算体系. 但是要强调的是: * 你的电脑里有合适的 GPU 显卡(NVIDIA), 且支持 CUDA 模块. 请在NVIDIA官网查询 * 必须安装 GPU 版的 Torch, 点击这里查看如何安装

用 GPU 训练 CNN¶

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

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

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

大功告成~

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

转移至 CPU¶

如果你有些计算还是需要在 CPU 上进行的话呢, 比如 plt 的可视化, 我们需要将这些计算或者数据转移至 CPU.