Pytorch07——半精度训练
GPU的性能主要分为两部分:算力和显存,前者决定了显卡计算的速度,后者则决定了显卡可以同时放入多少数据用于计算。在可以使用的显存数量一定的情况下,每次训练能够加载的数据更多(也就是batch size更大),则可以提高训练效率。另外有时候数据本身也比较大(比如3D图像、视频等),显存较小的情况下可能甚至batch size为1情况都无法实现,因此显存的大小十分重要。
我们观察Pytorch默认的浮点数存储方式用的是torch.float32,小数点后位数更多固然能够保证数据的精确性,但绝大多数场景其实并不需要那么精确,只保留一半的信息也不会影响结果,也就是使用torch.float16格式。由于数位减了一半,因此被称为半精度,具体如下图:
通过上图很明显的可以看到,使用半精度能够减少显存占用,使得显卡可以同时加载更多数据进行计算。
半精度训练的设置
在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)
...
本文来自博客园,作者:TCcjx,转载请注明原文链接:https://www.cnblogs.com/tccjx/articles/16026728.html