Automatic mixed precision for Pytorch 自动混合精度训练
参考的官方的链接 : Automatic mixed precision for Pytorch
混合精度训练:训练过程中将部分FP32转为FP16,从而实现训练速度的提升
快速入手混合精度训练
只需将下面的训练过程
... for input, target in data: optimizer.zero_grad() output = model(input) loss = loss_fn(output, target) loss.backward() optimizer.step()
改成
scaler = AmpScaler() ... for input, target in data: optimizer.zero_grad() # 前向过程 Auto-casting with autocast(): output = model(input) loss = loss_fn(output, target) # 反向传播 Gradient scaling scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()
就能实现训练速度上的提升，并且几乎不会造成精度上的损失。
这种方式还是比较可靠的,在yolov5上就使用的这种混合精度训练,并且达到了不错的效果。
