在Pytorch中使用GPU加速模型训练_python gpu加速

CPU->GPU

当我们需要训练一个复杂模型的时候，使用CPU来进行训练往往会花费非常多的时间。因此，一般情况下，我们都会选择在GPU (Graphics Process Unit)环境下进行训练，因为其支持高度的并行计算，可以有效地加快数据处理的速度，减少模型训练时间。Pytorch支持CUDA并行计算框架，该框架是由NVIDIA公司推出的基于自家GPU的并行计算框架。

CUDA 环境检验

如果我们希望在Pytorch实现GPU加速，首先我们需要确认GPU环境是否可用，我们可以使用如下命令来check。

torch.cuda.is_available()
1

返回结果是一个bool值，True或者False。如果是True, 代表GPU环境可用。

由于我自己的mac电脑上并没有安装GPU，一般我都会使用Google Colab中免费（限时）的GPU资源来训练参数量比较大的模型。

GPU训练的注意点

当我们确认GPU环境可用之后，接下来，我们需要将原本存放在CPU上的模型, 训练数据以及训练过程中使用到的额外的tensor都转移到GPU环境中, 示例代码如下。

# 设置选择哪个gpu
device='cuda:0'

#将写好的模型存放到GPU环境
model=model.to(device)

#将训练数据x,y转换到GPU环境
x,y=x.to(deivce),y.to(device)
1
2
3
4
5
6
7
8

除此之外，还有一些额外的tensor, 比如在使用nn.CrossEntropyLoss()时，我们想要对不同的class设置不同的权重， 我们需要向weight参数传入对应的权重张量，这个额外的张量也需要通过.to(device)转换到GPU环境中，这些细节要特别注意。

总之，在Pytorch中，对于数据转换这一块需要特别仔细，只要某个训练中需要用到的变量仍然存放在CPU上，Pytorch就会报错，说至少发现数据存放在cpu和cuda:0两个不同的环境下。

后续当我们需要将数据再次转回到CPU上时，我们可以使用.cpu()命令来实现。