Pytorch 使用GPU加速

68 阅读 0 评论 45 点赞

我是靠谱客的博主可爱小懒虫，最近开发中收集的这篇文章主要介绍Pytorch 使用GPU加速，觉得挺不错的，现在分享给大家，希望可以做个参考。

Pytorch中使用torch.device()选取并返回抽象出的设备，然后在定义的网络模块或者Tensor后面加上.to(device变量)就可以将它们搬到设备上了。

以上一篇代码为例，使用GPU设备：

device = torch.device('cuda:0') # 使用第一张显卡

需要将如下部分搬移到GPU上：

1. 定义的网络

net = MLP().to(device)

2. 损失函数

criteon = nn.CrossEntropyLoss()

3.每次取出的训练集和验证集的batch

data, target = data.to(device), target.to(device)

以上就是可爱小懒虫为你收集整理的Pytorch 使用GPU加速的全部内容，希望文章能够帮你解决Pytorch 使用GPU加速所遇到的程序开发问题。

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

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