今天pytorch 1.6正式发布了,本次更新的亮点在于引入了自动混合精度训练,详情见官网https://pytorch.org/blog/pytorch-1.6-released/
在此做一下简介
自动混合精度的意义在于加入了半精度的张量类型,这种类型可以在某些运算中具有更快的速度(如卷积和全连接层),官方文档中支持半精度的类型如下
__matmul__
, addbmm
, addmm
, addmv
, addr
, baddbmm
, bmm
, chain_matmul
,
conv1d
, conv2d
, conv3d
, conv_transpose1d
, conv_transpose2d
, conv_transpose3d
,
linear
, matmul
, mm
, mv
, prelu
此部分详细介绍见https://pytorch.org/docs/stable/amp.html#autocast-op-reference
有些暂时不支持半精度,比如LSTM和GRU,使用会报错,提示
RuntimeError: cuDNN error: CUDNN_STATUS_BAD_PARAM
【使用方法】
一个标准例子如下
from torch.cuda.amp import autocast as autocast, GradScaler # 创建model,默认是torch.FloatTensor model = Net().cuda() optimizer = optim.SGD(model.parameters(), ...) # 在训练最开始之前实例化一个GradScaler对象 scaler = GradScaler() for epoch in epochs: for input, target in data: optimizer.zero_grad() # 前向过程(model + loss)开启 autocast with autocast(): output = model(input) loss = loss_fn(output, target) # Scales loss,这是因为半精度的数值范围有限,因此需要用它放大 scaler.scale(loss).backward() # scaler.step() unscale之前放大后的梯度,但是scale太多可能出现inf或NaN # 故其会判断是否出现了inf/NaN # 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重, # 如果检测到出现了inf或者NaN,就跳过这次梯度更新,同时动态调整scaler的大小 scaler.step(optimizer) # 查看是否要更新scaler scaler.update()
更多实例以及Doc见官网
https://pytorch.org/docs/stable/notes/amp_examples.html#typical-mixed-precision-training
关于半精度的详细解释,推荐一篇知乎文章,见
PyTorch的自动混合精度(AMP) - Gemfield的文章 - 知乎 https://zhuanlan.zhihu.com/p/165152789
以上