自动混合精度示例 ¶

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

torch.autocast 的实例可以启用所选区域的自动转换。自动转换会自动选择操作精度以提升性能同时保持准确性。

torch.amp.GradScaler 帮助方便地执行梯度缩放的步骤。梯度缩放通过最小化梯度下溢，从而提高具有 float16 （默认在 CUDA 和 XPU 上）梯度的网络的收敛速度，具体解释如下。

torch.autocast 和 torch.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(device_type='cuda', dtype=torch.float16):
            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_() ）达到 $<=$ 某个用户设定的阈值。如果您尝试在不进行未缩放处理的情况下进行裁剪，梯度的范数/最大幅度也会被缩放，因此您请求的阈值（本意是未缩放梯度的阈值）将无效。

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

梯度裁剪

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

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast(device_type='cuda', dtype=torch.float16):
            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() 之前不要重复解除梯度缩放。

警告

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

与缩放梯度一起工作

梯度累积

梯度累积将梯度累加到一个有效批次的尺寸为 batch_per_iter * iters_to_accumulate （ * num_procs 如果是分布式的话）。缩放应该针对有效批次进行校准，这意味着进行无穷大/NaN 检查，如果发现无穷大/NaN 梯度则跳过步骤，并且缩放更新应该以有效批次粒度发生。此外，梯度应保持缩放，缩放因子应保持恒定，在累积给定有效批次的梯度时。如果梯度未缩放（或缩放因子发生变化）在累积完成之前，下一次反向传播将添加缩放梯度到未缩放梯度（或由不同因子缩放的梯度）之后，这将导致无法恢复累积的未缩放梯度 step 必须应用。

因此，如果您想 unscale_ 梯度（例如，允许未缩放的梯度裁剪），请在所有（缩放后的）即将到来的 step 梯度累积之后，调用 unscale_ 。另外，只有在迭代结束时调用 step 完整有效批次的情况下，才调用 update ：

scaler = GradScaler()

for epoch in epochs:
    for i, (input, target) in enumerate(data):
        with autocast(device_type='cuda', dtype=torch.float16):
            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() 的张量（s）进行缩放。因此，得到的梯度将进行缩放，在组合以创建惩罚值之前应进行反缩放。

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

下面是相同 L2 惩罚的示例：

scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast(device_type='cuda', dtype=torch.float16):
            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(device_type='cuda', dtype=torch.float16):
            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 应该只调用一次，在所有使用的优化器完成这一轮迭代之后：

scaler = torch.amp.GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer0.zero_grad()
        optimizer1.zero_grad()
        with autocast(device_type='cuda', dtype=torch.float16):
            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(device_type='cuda', dtype=torch.float16):
    # dp_model's internal threads will autocast.
    output = dp_model(input)
    # loss_fn also autocast
    loss = loss_fn(output)

分布式 DataParallel，每个进程一个 GPU ¶

该文档建议每个进程使用一个 GPU 以获得最佳性能。在这种情况下， DistributedDataParallel 不会内部生成线程，因此 autocast 和 GradScaler 的使用不受影响。

分布式数据并行，每个进程多个 GPU ¶

在这里， torch.nn.parallel.DistributedDataParallel 可能会为每个设备运行前向传递生成一个辅助线程，就像 torch.nn.DataParallel 一样。修复方法是相同的：将 autocast 作为模型 forward 方法的一部分应用，以确保在辅助线程中启用。

自动转换和自定义自动微分函数 ¶

如果您的网络使用自定义的自动微分函数（ torch.autograd.Function 的子类），则如果任何函数

• 接受多个浮点 Tensor 输入，

• 包装任何可自动转换的操作（参见自动转换操作参考），或

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

在所有情况下，如果你正在导入函数且无法修改其定义，一个安全的回退方案是在使用时禁用自动类型转换并强制在 float32 （或 dtype ）中执行，以避免错误发生：

with autocast(device_type='cuda', dtype=torch.float16):
    ...
    with autocast(device_type='cuda', dtype=torch.float16, enabled=False):
        output = imported_function(input1.float(), input2.float())

如果你是该函数的作者（或可以修改其定义），一个更好的解决方案是使用 torch.amp.custom_fwd() 和 torch.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 可以在任何地方调用，无需禁用 autocast 或手动转换输入：

mymm = MyMM.apply

with autocast(device_type='cuda', dtype=torch.float16):
    output = mymm(input1, input2)

需要特定 dtype 的函数 ¶

考虑一个需要 torch.float32 输入的自定义函数。将 custom_fwd(device_type='cuda', cast_inputs=torch.float32) 应用到 forward 和 custom_bwd(device_type='cuda') 应用到 backward 。如果 forward 在启用 autocast 的区域运行，装饰器会将浮点 Tensor 输入转换为 float32 并分配给指定的设备，例如本例中的 CUDA，并在 forward 和 backward 期间本地禁用 autocast：

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

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

func = MyFloat32Func.apply

with autocast(device_type='cuda', dtype=torch.float16):
    # func will run in float32, regardless of the surrounding autocast state
    output = func(input)