自动混合精度是一种能够提升训练效率的方法。它通过减少训练过程中的显存使用，从而提高 batch_size 大小，加快模型训练。
在 PyTorch 中张量默认使用的是 float32 类型，如果我们能够使用 float16 类型的话，那么将会很大程度上减少显存的使用，并且 float16 的计算效率要比 float32 要更好。
> 我们是否将默认使用的类型由 float32 设置为 float16，是否可以提升训练效率？
并不是这样的，float16 类型的缺点是对于非常小的浮点数会发生溢出问题，此时存储的值就可能会变成 0。
所以，对于整个计算过程中，有些计算适合使用 float16，而有些计算则适合使用 float32，即：对于模型训练来说，我们要根据情况来实现张量在 float16 和 float32 之间的切换。

有哪些计算可以使用 float16 ？

在 PyTorch 1.10 的文档中，可以转换为 float16 的操作：

__matmul__, addbmm, addmm, addmv, addr, baddbmm, bmm, chain_matmul, multi_dot, conv1d, conv2d, conv3d, conv_transpose1d, conv_transpose2d, conv_transpose3d, GRUCell, linear, LSTMCell, matmul, mm, mv, prelu, RNNCell

我们并不需要特别关心哪些操作会进行类型转换，这些由 PyTorch 框架来决定。使用自动混合精度训练，我们需要两个工具：

from torch.cuda.amp import autocast
from torch.cuda.amp import GradScaler

autocast 能够实现自动类型转换，将某些数据使用 float16 来表示，以提高计算效率，减少显存占用。以我们模型训练步骤为例：

1. 模型计算
2. 计算损失
3. 反向传播
4. 参数更新

模型计算和损失计算过程可以交由 autocast 来控制那些操作可以用 float16 来表示，那些可以用 float32 来表示。

当进行到反向传播时，会计算参数的梯度，有些梯度会很小，使得 float16 类型无法很好的存储该数据，所以这里使用 GradScaler 来 scaling 梯度，反向传播时参数的梯度会变得大一些，避免 float16 的数值溢出问题。
但是，进行参数更新时，会将梯度 unscaling 来更新参数。

即：正向计算时，我们使用 autocast 来自动实现类型转换，提高计算和减少显存使用，反向传播时使用 GradScaler 来避免数值溢出。

简单再次描述下 amp 的流程：
1. 开始 amp 之前，先将 float32 的参数拷贝一份；
2. 开始正向计算时，将 float32 的参数再拷贝一份，并将其类型转换为 float16；
3. 正向计算结束之后，将损失值乘以一个缩放因子再进行反向传播，目的是将所有的梯度值放大，避免数值太小而导致的数值溢出问题；
4. 参数更新之前，再将梯度值除以缩放因子，回归原本的梯度值再进行参数更新；
5. 最后更新下缩放因子。

1. 张量默认类型操作

我们可以使用 get_default_dtype 和 set_default_dtype 来获得和设置默认的张量类型。

from torch.cuda.amp import autocast
from torch.cuda.amp import GradScaler
import torch
import torch.nn as nn


def test01():

    # 获得张量默认类型
    print('默认张量类型:', torch.get_default_dtype())
    print('新创建张量类型:', torch.tensor(3.14).dtype)

    # 设置张量默认类型
    torch.set_default_dtype(torch.float16)
    print('新创建张量类型:', torch.tensor(3.14).dtype)


if __name__ == '__main__':
    test01()

程序执行结果：

默认张量类型: torch.float32
新创建张量类型: torch.float32
新创建张量类型: torch.float16

2. autocast

在 autocast 范围内时，PyTorch 会根据计算来自动进行类型转换，例如：

from torch.cuda.amp import autocast
import torch
import torch.nn as nn


def test02():

    a = torch.rand((8, 8), device="cuda")
    b = torch.rand((8, 8), device="cuda")
    print(a.dtype, b.dtype)

    # 没有使用 autocast 时，
    c = torch.mm(a, b)
    print('torch.mm(a, b):', c.dtype)

    # 使用 autocast 时，
    with autocast():

        # 会改变类型
        c = torch.mm(a, b)
        print('autocast torch.mm(a, b):', c.dtype)

        # 并不会改变类型
        d = torch.add(a, b)
        print('autocast torch.add(a, b):', d.dtype)



if __name__ == '__main__':
    test02()

程序执行结果：

torch.float32 torch.float32
torch.mm(a, b): torch.float32
autocast torch.mm(a, b): torch.float16
autocast torch.add(a, b): torch.float32

3. GradScaler

GradScaler 主要用于解决反向传播时，数值溢出的问题。解决思路，就是将梯度值乘以一个较大的数，将其数值表示到半精度类型能够存储的范围内，避免数值溢出问题。这里有个点需要注意，梯度值是一直变化的，一般随着训练的进行，误差越来越小，那么梯度值（也可以理解为参数的调整幅度）将会越来越小，那么乘以的缩放系数也要相应的变大。所以，GradScaler 中对梯度乘以的缩放系数是变化，会逐渐变大。

有些同学就会想，一直变大不会出现问题吗？可能也会出现数值太大，此时我们可以适当降低梯度的缩放系数。这个是由 GradScaler 内部做的判断。

1. scaler = GradScaler() 初始化梯度缩放器
2. scaler.scale(result) 将 result 乘以缩放因子
3. scaler.step(optimizer) 如果梯度中出现 nan 或者 inf 则不更新参数，否则更新参数
4. scaler.update() 缩放因子并不是固定不变，而是会更新。例如：发现 nan 或则 inf 则降低缩放因子，达到指定的迭代次数则提高缩放因子。

from torch.cuda.amp import autocast
from torch.cuda.amp import GradScaler
import torch
import torch.nn as nn
import torch.optim as optim


def test03():

    # 定义参数
    param = torch.tensor(2.0, requires_grad=True, device='cuda')
    # 优化器
    optimizer = optim.SGD([param], lr=0.1, momentum=0)

    # 计算过程
    result = param ** 2 + 3
    print(result)
    # 反向传播
    result.backward()
    print(param.grad)
    # 参数更新
    optimizer.step()
    print(param)


def test04():

    scaler = GradScaler()

    # 定义参数
    param = torch.tensor(2.0, requires_grad=True, device='cuda')
    # 优化器
    optimizer = optim.SGD([param], lr=0.1, momentum=0)

    # 计算过程
    result = param ** 2 + 3
    print(result)
    # scaling 之后反向传播，梯度值相当于扩大了 scale 倍
    scaler.scale(result).backward()
    print('缩放因子:', scaler.get_scale(), '梯度值:', param.grad)

    # 参数更新处理
    # optimizer.step()
    # step 函数尝试更新参数，如果发现梯度是 nan 或者 inf 则不进行参数更新
    scaler.step(optimizer)
    print(param)

    # 更新缩放因子，损失一直在变化，所以会根据实际情况来调整因子
    scaler.update()


if __name__ == '__main__':
    test03()
    print('-' * 50)
    test04()

程序执行结果：

tensor(7., device='cuda:0', grad_fn=<AddBackward0>)
tensor(4., device='cuda:0')
tensor(1.6000, device='cuda:0', requires_grad=True)
--------------------------------------------------
tensor(7., device='cuda:0', grad_fn=<AddBackward0>)
缩放因子: 65536.0 梯度值: tensor(262144., device='cuda:0')
tensor(1.6000, device='cuda:0', requires_grad=True)
65536.0