2.pytorch--混合精度训练

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

前言

提示：这里可以添加本文要记录的大概内容：

例如：随着人工智能的不断发展，机器学习这门技术也越来越重要，很多人都开启了学习机器学习，本文就介绍了机器学习的基础内容。

---

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

一、混合精度训练

原理：

将FP32的数据存储成FP16显然可以减小存储空间也可加快运算速度。但对于 FP16 不可避免的问题就是：表示的范围较窄，如下图所示，大量非 0 梯度会遇到溢出问题。

解决办法是：

对梯度乘一个 [公式] 的系数，称为 scale factor，把梯度 shift 到 FP16 的表示范围。而这一切，自动混合精度训练都给我们安排好了。

步骤：

GradScaler的工作就是在反向传播前给 loss 乘一个 scale factor，所以之后反向传播得到的梯度都乘了相同的 scale factor。并且为了不影响学习率，在梯度更新前将梯度unscale。总结amp的基本训练流程：

1. 维护一个 FP32 数值精度模型的副本
2. 在每个iteration
   1.拷贝并且转换成 FP16 模型
   2.前向传播（FP16 的模型参数）
   3.loss 乘 scale factor s
   4.反向传播（FP16 的模型参数和参数梯度）
   5.参数梯度乘 1/s
   6.利用 FP16 的梯度更新 FP32 的模型参数

---

二、使用步骤

1.引入库

代码如下（示例）：

import torch

2.使用混合精度训练

代码如下（示例）：

# amp依赖Tensor core架构，所以model参数必须是cuda tensor类型
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# GradScaler对象用来自动做梯度缩放
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        # 在autocast enable 区域运行forward
        with autocast():
            # model做一个FP16的副本，forward
            output = model(input)
            loss = loss_fn(output, target)
        # 用scaler，scale loss(FP16)，backward得到scaled的梯度(FP16)
        scaler.scale(loss).backward()
        # scaler 更新参数，会先自动unscale梯度
        # 如果有nan或inf，自动跳过
        scaler.step(optimizer)
        # scaler factor更新
        scaler.update()

---

总结

本文仅仅简单介绍了混合精度训练的使用，混合精度训练能够在现有硬件的情况下，为我们节省内存并提升模型的运算速度。