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

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

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

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

1. 先定义一个全局变量device：

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

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

2. 模型初始化时加入GPU：

 model = Model().to(device)

3. 模型计算时batch_x加入GPU，返回值转回cpu：

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

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

4. output使用完之后删除，以释放GPU显存：

   del output 

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