评价此页
★ ★ ★ ★ ★
recipes/recipes/amp_recipe

注意

跳转至页面底部下载完整示例代码。

自动混合精度#

创建日期：2020年9月15日 | 最后更新：2025年1月30日 | 最后验证：2024年11月5日

作者: Michael Carilli

torch.cuda.amp 为混合精度训练提供了便捷的方法，其中部分运算使用 torch.float32 (float) 数据类型，而其他运算则使用 torch.float16 (half)。线性层和卷积层等某些算子在 float16 或 bfloat16 下运行速度要快得多。而归约（reductions）等其他算子通常需要 float32 的动态范围。混合精度尝试为每个算子匹配最合适的数据类型，从而降低网络的运行时间和内存占用。

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

本教程（Recipe）首先测量一个简单网络在默认精度下的性能，然后逐步介绍如何添加 autocast 和 GradScaler，以在混合精度下运行相同的网络并提升性能。

您可以下载并运行本教程作为独立的 Python 脚本。唯一的要求是 PyTorch 1.6 或更高版本以及支持 CUDA 的 GPU。

混合精度主要使支持 Tensor Core 的架构（Volta、Turing、Ampere）受益。本教程在这些架构上应该会显示显著（2-3 倍）的加速。在较早的架构（Kepler、Maxwell、Pascal）上，您可能会观察到较小的加速。运行 nvidia-smi 可查看您 GPU 的架构。

import torch, time, gc

# Timing utilities
start_time = None

def start_timer():
    global start_time
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.synchronize()
    start_time = time.time()

def end_timer_and_print(local_msg):
    torch.cuda.synchronize()
    end_time = time.time()
    print("\n" + local_msg)
    print("Total execution time = {:.3f} sec".format(end_time - start_time))
    print("Max memory used by tensors = {} bytes".format(torch.cuda.max_memory_allocated()))

一个简单的网络#

以下线性层和 ReLU 的序列在使用混合精度时应能显示出加速效果。

def make_model(in_size, out_size, num_layers):
    layers = []
    for _ in range(num_layers - 1):
        layers.append(torch.nn.Linear(in_size, in_size))
        layers.append(torch.nn.ReLU())
    layers.append(torch.nn.Linear(in_size, out_size))
    return torch.nn.Sequential(*tuple(layers)).cuda()

batch_size、in_size、out_size 和 num_layers 的选择旨在足够大，以使 GPU 负载饱和。通常，当 GPU 饱和时，混合精度提供的加速效果最为明显。小型网络可能会受限于 CPU，在这种情况下，混合精度不会提高性能。尺寸的选择也确保了线性层参与计算的维度是 8 的倍数，从而允许在支持 Tensor Core 的 GPU 上使用 Tensor Core（参见下方的 故障排查）。

练习：改变参与计算的尺寸，观察混合精度的加速效果如何变化。 

batch_size = 512 # Try, for example, 128, 256, 513.
in_size = 4096
out_size = 4096
num_layers = 3
num_batches = 50
epochs = 3

device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.set_default_device(device)

# Creates data in default precision.
# The same data is used for both default and mixed precision trials below.
# You don't need to manually change inputs' ``dtype`` when enabling mixed precision.
data = [torch.randn(batch_size, in_size) for _ in range(num_batches)]
targets = [torch.randn(batch_size, out_size) for _ in range(num_batches)]

loss_fn = torch.nn.MSELoss().cuda()

默认精度#

如果不使用 torch.cuda.amp，以下简单网络会以默认精度（torch.float32）执行所有算子。 

net = make_model(in_size, out_size, num_layers)
opt = torch.optim.SGD(net.parameters(), lr=0.001)

start_timer()
for epoch in range(epochs):
    for input, target in zip(data, targets):
        output = net(input)
        loss = loss_fn(output, target)
        loss.backward()
        opt.step()
        opt.zero_grad() # set_to_none=True here can modestly improve performance
end_timer_and_print("Default precision:")

添加 torch.autocast#

torch.autocast 的实例充当上下文管理器，允许脚本中的特定区域以混合精度运行。

在这些区域内，CUDA 算子会以 autocast 选择的 dtype 运行，旨在在保持精度的同时提高性能。有关 autocast 为每个算子选择何种精度以及在何种情况下选择的详细信息，请参阅 Autocast 算子参考。 

for epoch in range(0): # 0 epochs, this section is for illustration only
    for input, target in zip(data, targets):
        # Runs the forward pass under ``autocast``.
        with torch.autocast(device_type=device, dtype=torch.float16):
            output = net(input)
            # output is float16 because linear layers ``autocast`` to float16.
            assert output.dtype is torch.float16

            loss = loss_fn(output, target)
            # loss is float32 because ``mse_loss`` layers ``autocast`` to float32.
            assert loss.dtype is torch.float32

        # Exits ``autocast`` before backward().
        # Backward passes under ``autocast`` are not recommended.
        # Backward ops run in the same ``dtype`` ``autocast`` chose for corresponding forward ops.
        loss.backward()
        opt.step()
        opt.zero_grad() # set_to_none=True here can modestly improve performance

添加 GradScaler#

梯度缩放 (Gradient scaling) 有助于防止在进行混合精度训练时，小幅度的梯度因数值过小而归零（“下溢”）。

torch.cuda.amp.GradScaler 可以方便地执行梯度缩放步骤。 

# Constructs a ``scaler`` once, at the beginning of the convergence run, using default arguments.
# If your network fails to converge with default ``GradScaler`` arguments, please file an issue.
# The same ``GradScaler`` instance should be used for the entire convergence run.
# If you perform multiple convergence runs in the same script, each run should use
# a dedicated fresh ``GradScaler`` instance. ``GradScaler`` instances are lightweight.
scaler = torch.amp.GradScaler("cuda")

for epoch in range(0): # 0 epochs, this section is for illustration only
    for input, target in zip(data, targets):
        with torch.autocast(device_type=device, dtype=torch.float16):
            output = net(input)
            loss = loss_fn(output, target)

        # Scales loss. Calls ``backward()`` on scaled loss to create scaled gradients.
        scaler.scale(loss).backward()

        # ``scaler.step()`` first unscales the gradients of the optimizer's assigned parameters.
        # If these gradients do not contain ``inf``s or ``NaN``s, optimizer.step() is then called,
        # otherwise, optimizer.step() is skipped.
        scaler.step(opt)

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

        opt.zero_grad() # set_to_none=True here can modestly improve performance

总结：“自动混合精度”#

（以下代码还展示了 enabled 参数，这是 autocast 和 GradScaler 的一个可选便利参数。如果设为 False，则 autocast 和 GradScaler 的调用将变为无操作 (no-op)。这允许在默认精度和混合精度之间切换，而无需使用 if/else 语句。）

use_amp = True

net = make_model(in_size, out_size, num_layers)
opt = torch.optim.SGD(net.parameters(), lr=0.001)
scaler = torch.amp.GradScaler("cuda" ,enabled=use_amp)

start_timer()
for epoch in range(epochs):
    for input, target in zip(data, targets):
        with torch.autocast(device_type=device, dtype=torch.float16, enabled=use_amp):
            output = net(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()
        scaler.step(opt)
        scaler.update()
        opt.zero_grad() # set_to_none=True here can modestly improve performance
end_timer_and_print("Mixed precision:")

检查/修改梯度（例如裁剪）#

由 scaler.scale(loss).backward() 产生的所有梯度都是经过缩放的。如果您希望在 backward() 和 scaler.step(optimizer) 之间修改或检查参数的 .grad 属性，您应该先使用 scaler.unscale_(optimizer) 对其进行反缩放。 

for epoch in range(0): # 0 epochs, this section is for illustration only
    for input, target in zip(data, targets):
        with torch.autocast(device_type=device, dtype=torch.float16):
            output = net(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()

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

        # Since the gradients of optimizer's assigned parameters are now unscaled, clips as usual.
        # You may use the same value for max_norm here as you would without gradient scaling.
        torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=0.1)

        scaler.step(opt)
        scaler.update()
        opt.zero_grad() # set_to_none=True here can modestly improve performance

保存/恢复#

要以位级精度保存/恢复启用了 Amp 的运行，请使用 scaler.state_dict 和 scaler.load_state_dict。

保存时，应将 scaler 的状态字典与常规的模型和优化器状态字典一起保存。请在迭代开始时的前向传播之前，或迭代结束后的 scaler.update() 之后进行此操作。 

checkpoint = {"model": net.state_dict(),
              "optimizer": opt.state_dict(),
              "scaler": scaler.state_dict()}
# Write checkpoint as desired, e.g.,
# torch.save(checkpoint, "filename")

恢复时，将 scaler 的状态字典与模型和优化器状态字典一起加载。根据需要读取检查点，例如 

dev = torch.cuda.current_device()
checkpoint = torch.load("filename",
                        map_location = lambda storage, loc: storage.cuda(dev))

net.load_state_dict(checkpoint["model"])
opt.load_state_dict(checkpoint["optimizer"])
scaler.load_state_dict(checkpoint["scaler"])

如果检查点是来自 没有 使用 Amp 的运行，而您想 使用 Amp 恢复训练，请像往常一样从检查点加载模型和优化器状态。因为检查点中不包含保存的 scaler 状态，所以请创建一个新的 GradScaler 实例。

如果检查点是来自 使用 Amp 的运行，而您想 不使用 Amp 恢复训练，请像往常一样从检查点加载模型和优化器状态，并忽略已保存的 scaler 状态。

推理/评估#

autocast 可以单独用于封装推理或评估的前向传播过程。无需 GradScaler。

进阶话题#

请参阅 自动混合精度示例 以获取进阶用例，包括：

  • 梯度累积 (Gradient accumulation)

  • 梯度惩罚/双向反向传播 (Gradient penalty/double backward)

  • 具有多个模型、优化器或损失的网络

  • 多 GPU (torch.nn.DataParallel 或 torch.nn.parallel.DistributedDataParallel)

  • 自定义 autograd 函数 (torch.autograd.Function 的子类)

如果您在同一脚本中进行多次收敛运行，则每次运行都应使用专门的全新 GradScaler 实例。GradScaler 实例非常轻量。

如果您要向调度器注册自定义 C++ 算子，请参阅调度器教程中的 autocast 章节。

故障排除#

使用 Amp 加速不明显#

  1. 您的网络可能未能使 GPU 负载饱和，因此受限于 CPU。Amp 对 GPU 性能的影响将无关紧要。

    • 使 GPU 饱和的一个经验法则是，在不导致内存溢出 (OOM) 的前提下，尽可能增加批次大小和/或网络尺寸。

    • 尽量避免过多的 CPU-GPU 同步（例如调用 .item()，或打印 CUDA 张量的值）。

    • 尽量避免连续进行大量小型的 CUDA 算子操作（如果可能，将它们合并为几个大型的 CUDA 算子）。

  2. 您的网络可能是 GPU 计算密集型（大量的 matmul/卷积），但您的 GPU 没有 Tensor Cores。在这种情况下，预期的加速效果会减小。

  3. matmul 的维度对 Tensor Core 不友好。确保 matmul 参与计算的尺寸是 8 的倍数。（对于带有编码器/解码器的 NLP 模型，这一点可能很微妙。此外，卷积过去对于 Tensor Core 的使用也有类似的尺寸限制，但对于 CuDNN 7.3 及更高版本，不存在此类限制。请参阅此处以获取指导。）

损失为 inf/NaN#

首先，检查您的网络是否符合 进阶用例。另请参阅 Prefer binary_cross_entropy_with_logits over binary_cross_entropy。

如果您确信您的 Amp 使用方法正确，则可能需要提交 Issue，但在执行此操作之前，收集以下信息会很有帮助：

  1. 分别禁用 autocast 或 GradScaler（通过向其构造函数传递 enabled=False），查看 infs/NaNs 是否持续存在。

  2. 如果您怀疑网络的一部分（例如复杂的损失函数）发生溢出，请在该前向传播区域以 float32 运行，查看 infs/NaNs 是否持续存在。Autocast 文档字符串 <https://pytorch.ac.cn/docs/stable/amp.html#torch.autocast>`_ 的最后一个代码片段展示了如何强制子区域在 ``float32 下运行（通过局部禁用 autocast 并对该子区域的输入进行强制类型转换）。

类型不匹配错误（可能表现为 CUDNN_STATUS_BAD_PARAM）#

Autocast 尝试涵盖所有受益于或需要转换的算子。受明确支持的算子的选择是基于数值属性以及工程经验。如果您在启用了 autocast 的前向区域或随后的反向传播中看到类型不匹配错误，则可能是 autocast 漏掉了一个算子。

请提交包含错误回溯的 Issue。在运行脚本之前，请设置 export TORCH_SHOW_CPP_STACKTRACES=1，以提供有关哪个后端算子失败的精细信息。

在此页面上
PyTorch 库