CUDA 自动混合精度示例

通常，“自动混合精度训练”是指同时使用 torch.cuda.amp.autocast 和 torch.cuda.amp.GradScaler 进行训练。

Instances of torch.cuda.amp.autocast 使选定区域启用自动类型转换。 自动类型转换会自动为GPU操作选择精度以提高性能 同时保持准确性。

Instances of torch.cuda.amp.GradScaler 帮助方便地执行梯度缩放的步骤。梯度缩放通过最小化梯度下溢来提高具有 float16 梯度的网络的收敛性，如 此处 所解释。

torch.cuda.amp.autocast 和 torch.cuda.amp.GradScaler 是模块化的。 在下面的示例中，每个都按照其单独的文档说明使用。

（此处的示例仅供参考。请参阅 自动混合精度配方 以获取可运行的演练。）

典型的混合精度训练

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Creates a GradScaler once at the beginning of training.
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # Runs the forward pass with autocasting.
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
        # Backward passes under autocast are not recommended.
        # Backward ops run in the same dtype autocast chose for corresponding forward ops.
        scaler.scale(loss).backward()

        # scaler.step() first unscales the gradients of the optimizer's assigned params.
        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,
        # otherwise, optimizer.step() is skipped.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        scaler.update()

处理未缩放的梯度

所有由scaler.scale(loss).backward()生成的梯度都被缩放了。如果你想在backward()和scaler.step(optimizer)之间修改或检查参数的.grad属性，你应该首先取消缩放它们。例如，梯度裁剪会调整一组梯度，使它们的全局范数（见torch.nn.utils.clip_grad_norm_()）或最大幅度（见torch.nn.utils.clip_grad_value_()）达到某个用户设定的阈值$<=$。如果你尝试在不取消缩放的情况下进行裁剪，梯度的范数/最大幅度也会被缩放，因此你请求的阈值（本应是针对未缩放梯度的阈值）将无效。

scaler.unscale_(optimizer) 对 optimizer 分配的参数持有的梯度进行缩放。 如果你的模型或其他模型包含分配给另一个优化器的其他参数 （比如 optimizer2），你也可以单独调用 scaler.unscale_(optimizer2) 来对这些 参数的梯度进行缩放。

梯度裁剪

在裁剪之前调用 scaler.unscale_(optimizer) 可以让你像往常一样裁剪未缩放的梯度：

scaler = GradScaler()

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

        # Unscales the gradients of optimizer's assigned params in-place
        scaler.unscale_(optimizer)

        # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        # optimizer's gradients are already unscaled, so scaler.step does not unscale them,
        # although it still skips optimizer.step() if the gradients contain infs or NaNs.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        scaler.update()

scaler 条记录表明 scaler.unscale_(optimizer) 在本次迭代中已经为该优化器调用过，因此 scaler.step(optimizer) 知道在（内部）调用 optimizer.step() 之前不需要重复地取消缩放梯度。

警告

unscale_ 应该在每个优化器的每次 step 调用中仅调用一次， 并且仅在为该优化器分配的所有参数都累积了梯度之后。 对于给定的优化器，在每次 step 之间两次调用 unscale_ 将触发 RuntimeError。

处理缩放梯度

梯度累积

梯度累积会在一个有效批次大小为 batch_per_iter * iters_to_accumulate 的情况下累加梯度 （如果使用分布式则为 * num_procs）。应针对有效批次进行缩放校准，这意味着需要检查 inf/NaN 值， 在发现 inf/NaN 梯度时跳过步骤，并且缩放更新应在有效批次粒度上进行。此外，在累积某个有效批次的梯度期间， 梯度应保持缩放状态，缩放因子也应保持不变。如果在累积完成之前梯度被反缩放（或缩放因子发生变化）， 下一次反向传播将把缩放后的梯度添加到未缩放的梯度中（或用不同因子缩放的梯度），之后就无法恢复累积的未缩放梯度 step 必须应用。

因此，如果你想获取 unscale_ 梯度（例如，允许裁剪未缩放的梯度）， 请在 unscale_ 之前调用 step，在所有（缩放后的）梯度为即将进行的 step 累积之后。此外，仅在你调用了 update 的迭代结束时调用 step 以获得完整的有效批次：

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)
            loss = loss / iters_to_accumulate

        # Accumulates scaled gradients.
        scaler.scale(loss).backward()

        if (i + 1) % iters_to_accumulate == 0:
            # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)

            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

梯度惩罚

梯度惩罚的实现通常使用 torch.autograd.grad() 创建梯度，将它们组合起来创建惩罚值， 并将惩罚值添加到损失中。

这是一个普通的L2惩罚示例，不包含梯度缩放或自动类型转换：

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

        # Creates gradients
        grad_params = torch.autograd.grad(outputs=loss,
                                          inputs=model.parameters(),
                                          create_graph=True)

        # Computes the penalty term and adds it to the loss
        grad_norm = 0
        for grad in grad_params:
            grad_norm += grad.pow(2).sum()
        grad_norm = grad_norm.sqrt()
        loss = loss + grad_norm

        loss.backward()

        # clip gradients here, if desired

        optimizer.step()

要实现带有梯度缩放的梯度惩罚，传递给torch.autograd.grad()的outputs个张量应进行缩放。因此，生成的梯度也会被缩放，在组合以创建惩罚值之前，应该先将其反向缩放。

此外，惩罚项的计算是前向传播的一部分，因此应该位于 autocast 上下文中。

这就是相同L2惩罚的样子：

scaler = GradScaler()

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

        # Scales the loss for autograd.grad's backward pass, producing scaled_grad_params
        scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss),
                                                 inputs=model.parameters(),
                                                 create_graph=True)

        # Creates unscaled grad_params before computing the penalty. scaled_grad_params are
        # not owned by any optimizer, so ordinary division is used instead of scaler.unscale_:
        inv_scale = 1./scaler.get_scale()
        grad_params = [p * inv_scale for p in scaled_grad_params]

        # Computes the penalty term and adds it to the loss
        with autocast():
            grad_norm = 0
            for grad in grad_params:
                grad_norm += grad.pow(2).sum()
            grad_norm = grad_norm.sqrt()
            loss = loss + grad_norm

        # Applies scaling to the backward call as usual.
        # Accumulates leaf gradients that are correctly scaled.
        scaler.scale(loss).backward()

        # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)

        # step() and update() proceed as usual.
        scaler.step(optimizer)
        scaler.update()

处理多个模型、损失函数和优化器

如果你的网络有多个损失，你必须分别在每个损失上单独调用 scaler.scale. 如果你的网络有多个优化器，你可以分别在任意一个优化器上单独调用 scaler.unscale_, 并且你必须分别在每个优化器上单独调用 scaler.step。

但是，scaler.update 只应在所有在此迭代中使用的优化器都已执行 step 之后调用一次：

scaler = torch.cuda.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer0.zero_grad()
        optimizer1.zero_grad()
        with autocast():
            output0 = model0(input)
            output1 = model1(input)
            loss0 = loss_fn(2 * output0 + 3 * output1, target)
            loss1 = loss_fn(3 * output0 - 5 * output1, target)

        # (retain_graph here is unrelated to amp, it's present because in this
        # example, both backward() calls share some sections of graph.)
        scaler.scale(loss0).backward(retain_graph=True)
        scaler.scale(loss1).backward()

        # You can choose which optimizers receive explicit unscaling, if you
        # want to inspect or modify the gradients of the params they own.
        scaler.unscale_(optimizer0)

        scaler.step(optimizer0)
        scaler.step(optimizer1)

        scaler.update()

每个优化器都会检查其梯度是否存在无穷大或NaN值，并独立决定是否跳过该步骤。这可能导致一个优化器跳过该步骤，而另一个则不跳过。由于跳过步骤的情况很少发生（每几百次迭代一次），这不应妨碍收敛。如果您在向多优化器模型添加梯度缩放后观察到收敛效果不佳，请报告一个错误。

使用多个GPU进行工作

此处描述的问题仅影响 autocast。 GradScaler 的使用方式保持不变。

在单个进程中使用DataParallel

即使 torch.nn.DataParallel 生成线程在每个设备上运行前向传递。 自动转换状态会在每个线程中传播，以下代码将正常工作：

model = MyModel()
dp_model = nn.DataParallel(model)

# Sets autocast in the main thread
with autocast():
    # dp_model's internal threads will autocast.
    output = dp_model(input)
    # loss_fn also autocast
    loss = loss_fn(output)

分布式数据并行处理，每个进程一个GPU

torch.nn.parallel.DistributedDataParallel’s documentation recommends one GPU per process for best performance. In this case, DistributedDataParallel does not spawn threads internally, so usages of autocast and GradScaler are not affected.

分布式数据并行处理，每个进程使用多个GPU

这里 torch.nn.parallel.DistributedDataParallel 可能在每个设备上启动一个辅助线程来运行前向传播， 像 torch.nn.DataParallel 一样。 修复方法相同： 将自动转换应用为模型的 forward 方法的一部分，以确保在辅助线程中启用它。

自动混合精度和自定义自动梯度函数

如果您的网络使用 自定义 autograd 函数 （torch.autograd.Function 的子类），则需要进行更改以实现与自动转换的兼容性，如果任何函数

• 接受多个浮点Tensor输入，

• 封装任何可自动转换的操作（请参阅自动转换操作参考），或

• 需要特定的dtype（例如，如果它封装了 CUDA 扩展 这些扩展仅针对dtype进行了编译）。

在所有情况下，如果你正在导入该函数并且无法更改其定义，则一个安全的回退方法是在出现错误的任何使用点禁用自动转换并强制执行 float32（或 dtype）：

with autocast():
    ...
    with autocast(enabled=False):
        output = imported_function(input1.float(), input2.float())

如果你是函数的作者（或可以修改其定义），一个更好的解决方案是使用 torch.cuda.amp.custom_fwd() 和 torch.cuda.amp.custom_bwd() 装饰器，如下所示的相关情况。

具有多个输入或可自动转换操作的函数

应用 custom_fwd 和 custom_bwd（不带参数）到 forward 和 backward 分别。这些确保 forward 在当前的自动类型转换状态下执行，而 backward 在与 forward 相同的自动类型转换状态下执行（这可以防止类型不匹配错误）：

class MyMM(torch.autograd.Function):
    @staticmethod
    @custom_fwd
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        return a.mm(b)
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        a, b = ctx.saved_tensors
        return grad.mm(b.t()), a.t().mm(grad)

现在 MyMM 可以在任何地方调用，而无需禁用自动类型转换或手动转换输入：

mymm = MyMM.apply

with autocast():
    output = mymm(input1, input2)

需要特定dtype的函数

考虑一个需要 torch.float32 个输入的自定义函数。 将 custom_fwd(cast_inputs=torch.float32) 应用于 forward 并将 custom_bwd（无参数）应用于 backward。 如果 forward 在启用了自动类型转换的区域中运行，装饰器会将浮点 CUDA 张量 输入转换为 float32，并在 forward 和 backward 期间本地禁用自动类型转换：

class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        ...
        return fwd_output
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        ...

现在 MyFloat32Func 可以在任何地方调用，无需手动禁用自动转换或转换输入：

func = MyFloat32Func.apply

with autocast():
    # func will run in float32, regardless of the surrounding autocast state
    output = func(input)