Pytorch07——半精度训练

GPU的性能主要分为两部分：算力和显存，前者决定了显卡计算的速度，后者则决定了显卡可以同时放入多少数据用于计算。在可以使用的显存数量一定的情况下，每次训练能够加载的数据更多（也就是batch size更大）,则可以提高训练效率。另外有时候数据本身也比较大（比如3D图像、视频等），显存较小的情况下可能甚至batch size为1情况都无法实现，因此显存的大小十分重要。

我们观察Pytorch默认的浮点数存储方式用的是torch.float32，小数点后位数更多固然能够保证数据的精确性，但绝大多数场景其实并不需要那么精确，只保留一半的信息也不会影响结果，也就是使用torch.float16格式。由于数位减了一半，因此被称为半精度，具体如下图：

通过上图很明显的可以看到，使用半精度能够减少显存占用，使得显卡可以同时加载更多数据进行计算。

1|0半精度训练的设置

在Pytorch中使用autocast配置半精度训练，同时需要在下面三处加以设置：

import autocast

  
from torch.cuda.amp import autocast

模型设置

在模型定义中，使用python的装饰器方法，用autocast装饰模型中的forward函数。关于装饰器的使用，参考下面：

 @autocast()
 
def forward(self,x):
 
　　...
 
　　return x

训练过程

在训练过程中，只需要将数据输入模型及其之后的部分放入"with autocast():"即可：

  
for x in train_loader:
 
　　x = x.cuda()
 
　　with autocast():
 
　　output = model(x)
 
　　...

__EOF__

本文作者：DeepBrainBoy
本文链接：https://www.cnblogs.com/tccjx/articles/16026728.html
关于博主：评论和私信会在第一时间回复。或者直接私信我。
版权声明：本博客所有文章除特别声明外，均采用 BY-NC-SA 许可协议。转载请注明出处！
声援博主：如果您觉得文章对您有帮助，可以点击文章右下角【推荐】一下。您的鼓励是博主的最大动力！