pytorch混合精度计算代码

---

提示：以下是本篇文章正文内容，下面案例可供参考

一、混合精度计算

为了想让深度学习训练可以使用FP16的好处，又要避免精度溢出和舍入误差。于是可以通过FP16和FP32的混合精度训练（Mixed-Precision），混合精度训练过程中可以引入权重备份（Weight Backup）、损失放大（Loss Scaling）、精度累加（Precision Accumulated）三种相关的技术。

二、使用步骤

1.引入库

代码如下（示例）：

import torch

2.设置混合精度计算开关

代码如下（示例）：

fp16 = True#False

3.设置scalar

torch 1.2不支持amp，建议使用torch 1.7.1及以上正确使用fp16。此设置在开始训练步骤之前设置

if fp16:
        from torch.cuda.amp import GradScaler as GradScaler
        scaler = GradScaler()
else:
        scaler = None

4.在训练中设置混合精度计算

 		#   清零梯度
        #----------------------#
        optimizer.zero_grad()
        if not fp16:
            #----------------------#
            #   前向传播
            #----------------------#
            outputs     = model_train(images)
            #----------------------#
            #   计算损失
            #----------------------#
            loss_value  = nn.CrossEntropyLoss()(outputs, targets)

            loss_value.backward()
            optimizer.step()
        else:
            from torch.cuda.amp import autocast
            with autocast():
                #----------------------#
                #   前向传播
                #----------------------#
                outputs     = model_train(images)
                #----------------------#
                #   计算损失
                #----------------------#
                loss_value  = nn.CrossEntropyLoss()(outputs, targets)

            #----------------------#
            #   反向传播
            #----------------------#
            scaler.scale(loss_value).backward()
            scaler.step(optimizer)
            scaler.update()

---

总结

使用混合精度计算可以减少约一半的显存，但是要注意需要torch1.7.1以上。