Pytorch torch.cuda.amp 自动混合精度

自动混合精度

Pytorch的自动混合精度是由torch.cuda.amp模块带来的

from torch.cuda.amp import autocast as autocast

Pytorch的amp模块里面有两种精度的Tensor，torch.FloatTensor和torch.HalfTensor。

torch.cuda.amp只能在cuda上使用，这个功能正是NVIDIA的开发人员贡献到Pytorch项目中的。

而只有支持Tensor core的CUDA硬件才能享受到AMP的好处(例如:2080Ti显卡)。

Tensor Core是一种矩阵累加的计算单元，每个Tensor Core每个始终执行64个浮点混合精度 *** 作，同时降低一半的显存访问和存储。

一般我们的所说的自动混合精度就是在CUDA设备上使用torch.cuda.amp.autocast和torch.cuda.amp.GradScaler来进行训练。

为什么要混合精度进行计算呢？直接使用torch.FloatTensor计算不就行了吗？因为torch.HalfTensor的优势就是存储小，计算快并且可以更好的利用CUDA设备的Tensor Core进行计算加速。

torch.HalfTensor的劣势就是数值范围小(更容易overflow/underflow)、舍入误差(Rounding Error, 导致一些微笑的梯度信息达不到16bit精度的最低分辨率，从而丢失).

我们只是使用torch.HalfTensor具有优势的一面，为了消除torch.HalfTensor的劣势，主要采用了两种解决方案:
1: 梯度Scale, 即：torch.cuda.amp.GradScaler，通过放大loss值来防止梯度的underflow(主要是BP的时候计算梯度的时候使用，真正在更新权重的时候还是要把方法的梯度再unscale回去)
2: 回到torch.FloatTensor,这就是混合，在Pytorch1.6之后的AMP中

__matmul__
addbmm
addmm
addmv
addr
baddbmm
bmm
chain_matmul
conv1d
conv2d
conv3d
conv_transpose1d
conv_transpose2d
conv_transpose3d
linear
matmul
mm
mv
prelu

使用步骤：

1. autocast

from torch.cuda.amp import autocast as autocast
for input, target in data:
optimizer.zero_grad()
with autocast():
	output = model(input)
	loss = loss_fn(output, target)
loss.backward()
optimzier.step()

进入到autocast之后，模型里面的上面那些列出来的CUDA ops会把tensor的dtype转换为半精度浮点类型，从而在不损失训练精度的情况下加快运算。

刚进入autocast的上下文时,tensor可以是任何类型。

2.结合GradScaler来进行计算

from torch.cuda.amp import autocast as autocast

scaler = GradScaler()
optimizer.zero_grad()
with autocast():
	output = model(input)
	loss = loss_fn(output, target)
scaler.scale(loss).backward()
optimzier.step()
scaler.step(optimizer)
scaler.update()

一般情况下Scaler的大小在每次迭代过程中动态的估计，尽量的减少Rounding Error但是也不能太大，太大容易overflow（inf NaN）
如果出现了inf或者NaN,scaler.step(optimzier)会自动忽略这次更新