Pytorch 中的 torch.cuda.amp.autocast 和 GradScaler 联合使用实现 AMP

1 amp 模块的作用

在 pytorch 中, 自动混合精度可以在神经网络推理过程中，针对不同的层，采用不同的数据精度进行计算，从而实现节省显存和加快速度的目的, 那么如何实现 amp?

• 使用 autocast + GradScaler 两者相结合的方法

• autocast 使用半精度加速训练, 但是半精度加速训练的 troch.HalfTensor 有一些不足, 针对这些不足, 使用 GradScaler 来解决, 具体原因见另一篇和 GradeScaler 有关的文章

2 示例

2.1 单独使用 autocast

import torch
from torch.cuda.amp import autocast

# 创建一个 model,默认的数据类型是 torch.FloatTensor
model = Net().cuda()
optimizer = torch.cuda.optim.SGD(moeld.parameters(), ... )

for input, target in data:
    optimizer.zero_grad()
    
    # 对前向传播中的 model + loss 计算时使用 autocast
    with autocast():
        output  = model(input)
        loss = loss_fn(output, target)

    # 反向传播 BP 在 autocast 以外
    loss.backward()
    optimizer.step()

2.2 GradScaler 与 autocast 联合使用

两者结合实现 AMP 方法

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

model = Net().cuda()
optimizer = torch.optim.SGD(model.parameters(), ...)

# 在训练开始之前实例化一个 GradeScaler 对象
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # 在 autocast() 后使用 GradScaler()
        # Backpropagation
        scaler.scale(loss).backward()
        # update generator weights
        scaler.step(optimizer)
        scaler.update()

2.3 多卡训练 nn.DataParallel

单卡训练的话上面的代码已经够了, 在2080ti上能减少至少1/3的显存, 要是想多卡跑的话仅仅这样还不够，会发现在forward里面的每个结果都还是float32的

解决方法:

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()

    def forward(self, input_data):
        with autocast():
            # code
        return

把forward里面的代码用autocast代码块方式运行