【Pytorch】使用GPU加速模型的训练

63 阅读 0 评论 42 点赞

我是靠谱客的博主不安果汁，最近开发中收集的这篇文章主要介绍【Pytorch】使用GPU加速模型的训练，觉得挺不错的，现在分享给大家，希望可以做个参考。

为什么要用GPU加速，因为GPU有大量的并行计算单元，能够至少提升一个数量级的计算速度。
关于GPU加速的原理不赘述，反正很快就完事了，同一个项目我用服务器的CPU(E5-2650)需要半小时，用GPU(1080Ti)只需要200多秒，快得不是一点半点。

下面介绍如何在Pytorch中使用GPU加速，主要有4个点：

device = torch.device('cuda:0')

其中，cuda:0可以指定GPU编号，如果有多个GPU，可以在linux中输入 nvidia-smi 查看GPU信息

 model = Model().to(device)

output = model(batch_x.to(device)).cpu()

转回cpu是为了不影响后续计算


del output

需要注意的地方：千万别在一开始就把训练集和测试集放到GPU里，以免爆内存，只需要在前向传播时，参数后.to(device)即可，返回值记得转回cpu，其他的loss函数什么的不需要额外操作

以上就是不安果汁为你收集整理的【Pytorch】使用GPU加速模型的训练的全部内容，希望文章能够帮你解决【Pytorch】使用GPU加速模型的训练所遇到的程序开发问题。

如果觉得靠谱客网站的内容还不错，欢迎将靠谱客网站推荐给程序员好友。

本图文内容来源于网友提供，作为学习参考使用，或来自网络收集整理，版权属于原作者所有。