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

加油!

感谢!

努力!