PyTorch在GPU上训练模型

Posted 2021-08-29 ZSYL

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了PyTorch在GPU上训练模型相关的知识，希望对你有一定的参考价值。

为了真正利用Pytorch中Tensor的优秀属性, 加速模型的训练, 我们可以将训练过程转移到GPU上进行.

1. 首先要定义设备, 如果CUDA是可用的则被定义成GPU, 否则被定义成CPU.

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print(device)

输出结果:

cuda:0

2. 当训练模型的时候, 只需要将模型转移到GPU上, 同时将输入的图片和标签页转移到GPU上即可.

# 将模型转移到GPU上
net.to(device)

# 将输入的图片张量和标签张量转移到GPU上
inputs, labels = data[0].to(device), data[1].to(device)

前提是电脑上安装的PyTorch是GPU版本的。

安装PyTorch GPU版本，请参考：link

加油!

感谢!

努力!

以上是关于PyTorch在GPU上训练模型的主要内容，如果未能解决你的问题，请参考以下文章