自动混合精度包 - torch.amp

torch.amp 提供了混合精度的便利方法，其中一些操作使用 torch.float32 (float) 数据类型，而其他操作使用较低精度的浮点数据类型 (lower_precision_fp)：torch.float16 (half) 或 torch.bfloat16。一些算子，如线性层和卷积，在 lower_precision_fp 中速度快得多。其他算子，如归约（reductions），通常需要 float32 的动态范围。混合精度尝试将每个算子与其适当的数据类型匹配。

通常，使用 torch.float16 数据类型的“自动混合精度训练”会一起使用 torch.autocast 和 torch.amp.GradScaler，如自动混合精度示例和自动混合精度秘籍所示。然而，torch.autocast 和 torch.GradScaler 是模块化的，如果需要，可以单独使用。如torch.autocast 的 CPU 示例部分所示，在 CPU 上使用 torch.bfloat16 数据类型的“自动混合精度训练/推理”仅使用 torch.autocast。

警告

torch.cuda.amp.autocast(args...) 和 torch.cpu.amp.autocast(args...) 将被弃用。请改用 torch.autocast("cuda", args...) 或 torch.autocast("cpu", args...)。torch.cuda.amp.GradScaler(args...) 和 torch.cpu.amp.GradScaler(args...) 将被弃用。请改用 torch.GradScaler("cuda", args...) 或 torch.GradScaler("cpu", args...)。

torch.autocast 和 torch.cpu.amp.autocast 是 1.10 版本中的新增功能。

自动类型转换

torch.amp.autocast_mode.is_autocast_available(device_type)

返回一个布尔值，指示 device_type 上是否可用 autocast。

参数

device_type (str) – 要使用的设备类型。可能的值包括：‘cuda’、‘cpu’、‘mtia’、‘xpu’ 等等。类型与 torch.device 的 type 属性相同。因此，您可以使用 Tensor.device.type 获取张量的设备类型。

返回类型

bool

class torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None)

autocast 的实例可用作上下文管理器或装饰器，允许您的脚本中的区域以混合精度运行。

在这些区域中，算子以 autocast 选择的算子特定数据类型运行，以提高性能同时保持精度。详见Autocast 算子参考。

进入启用 autocast 的区域时，张量可以是任何类型。使用 autocasting 时，不应对模型或输入调用 half() 或 bfloat16()。

autocast 应该只包裹网络的正向传播，包括损失计算。不建议在 autocast 下进行反向传播。反向算子以 autocast 用于相应正向算子的相同类型运行。

CUDA 设备示例

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

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

    # Enables autocasting for the forward pass (model + loss)
    with torch.autocast(device_type="cuda"):
        output = model(input)
        loss = loss_fn(output, target)

    # Exits the context manager before backward()
    loss.backward()
    optimizer.step()

有关更复杂场景（例如，梯度惩罚、多个模型/损失、自定义 autograd 函数）中的使用（以及梯度缩放），请参阅自动混合精度示例和自动混合精度秘籍。

autocast 也可用作装饰器，例如，用于模型的 forward 方法

class AutocastModel(nn.Module):
    ...
    @torch.autocast(device_type="cuda")
    def forward(self, input):
        ...

在启用 autocast 的区域中生成的浮点张量可能是 float16。返回到禁用 autocast 的区域后，将它们与不同数据类型的浮点张量一起使用可能会导致类型不匹配错误。如果发生这种情况，请将 autocast 区域中生成的张量转换回 float32（如果需要，也可以转换为其他数据类型）。如果来自 autocast 区域的张量已经是 float32，则转换是空操作，不会产生额外开销。CUDA 示例

# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")

with torch.autocast(device_type="cuda"):
    # torch.mm is on autocast's list of ops that should run in float16.
    # Inputs are float32, but the op runs in float16 and produces float16 output.
    # No manual casts are required.
    e_float16 = torch.mm(a_float32, b_float32)
    # Also handles mixed input types
    f_float16 = torch.mm(d_float32, e_float16)

# After exiting autocast, calls f_float16.float() to use with d_float32
g_float32 = torch.mm(d_float32, f_float16.float())

CPU 训练示例

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

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

        # Runs the forward pass with autocasting.
        with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
            output = model(input)
            loss = loss_fn(output, target)

        loss.backward()
        optimizer.step()

CPU 推理示例

# Creates model in default precision
model = Net().eval()

with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    for input in data:
        # Runs the forward pass with autocasting.
        output = model(input)

带 Jit Trace 的 CPU 推理示例

class TestModel(nn.Module):
    def __init__(self, input_size, num_classes):
        super().__init__()
        self.fc1 = nn.Linear(input_size, num_classes)
    def forward(self, x):
        return self.fc1(x)

input_size = 2
num_classes = 2
model = TestModel(input_size, num_classes).eval()

# For now, we suggest to disable the Jit Autocast Pass,
# As the issue: https://github.com/pytorch/pytorch/issues/75956
torch._C._jit_set_autocast_mode(False)

with torch.cpu.amp.autocast(cache_enabled=False):
    model = torch.jit.trace(model, torch.randn(1, input_size))
model = torch.jit.freeze(model)
# Models Run
for _ in range(3):
    model(torch.randn(1, input_size))

启用 autocast 的区域中的类型不匹配错误是 bug；如果您观察到这种情况，请提交问题。

autocast(enabled=False) 子区域可以嵌套在启用 autocast 的区域中。局部禁用 autocast 可能有用，例如，如果您想强制子区域以特定 dtype 运行。禁用 autocast 使您可以显式控制执行类型。在子区域中，来自周围区域的输入在使用前应转换为 dtype

# Creates some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cuda")
b_float32 = torch.rand((8, 8), device="cuda")
c_float32 = torch.rand((8, 8), device="cuda")
d_float32 = torch.rand((8, 8), device="cuda")

with torch.autocast(device_type="cuda"):
    e_float16 = torch.mm(a_float32, b_float32)
    with torch.autocast(device_type="cuda", enabled=False):
        # Calls e_float16.float() to ensure float32 execution
        # (necessary because e_float16 was created in an autocasted region)
        f_float32 = torch.mm(c_float32, e_float16.float())

    # No manual casts are required when re-entering the autocast-enabled region.
    # torch.mm again runs in float16 and produces float16 output, regardless of input types.
    g_float16 = torch.mm(d_float32, f_float32)

autocast 状态是线程本地的。如果想在新线程中启用它，必须在该线程中调用上下文管理器或装饰器。当每个进程使用多个 GPU 时，这会影响 torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel（详见使用多个 GPU）。

参数

• device_type (str, required) – 要使用的设备类型。可能的值包括：‘cuda’、‘cpu’、‘mtia’、‘xpu’ 和 ‘hpu’。类型与 torch.device 的 type 属性相同。因此，您可以使用 Tensor.device.type 获取张量的设备类型。

• enabled (bool, optional) – 该区域是否应启用 autocasting。默认值: True

• dtype (torch_dtype, optional) – autocast 中运行的算子使用的数据类型。如果 dtype 为 None，则使用 get_autocast_dtype() 给出的默认值（CUDA 为 torch.float16，CPU 为 torch.bfloat16）。默认值: None

• cache_enabled (bool, optional) – autocast 内部的权重缓存是否应启用。默认值: True

torch.amp.custom_fwd(fwd=None, *, device_type, cast_inputs=None)

为自定义 autograd 函数的 forward 方法创建辅助装饰器。

Autograd 函数是 torch.autograd.Function 的子类。详见示例页面。

参数

• device_type (str) – 要使用的设备类型。‘cuda’、‘cpu’、‘mtia’、‘xpu’ 等等。类型与 torch.device 的 type 属性相同。因此，您可以使用 Tensor.device.type 获取张量的设备类型。

• cast_inputs (torch.dtype or None, optional, default=None) – 如果不是 None，当 forward 在启用 autocast 的区域中运行时，会将输入的浮点张量转换为目标数据类型（非浮点张量不受影响），然后在禁用 autocast 的情况下执行 forward。如果为 None，则 forward 的内部算子将以当前的 autocast 状态执行。

注意

如果在启用 autocast 的区域外部调用被装饰的 forward，则 custom_fwd 是空操作，且 cast_inputs 无效。

torch.amp.custom_bwd(bwd=None, *, device_type)

为自定义 autograd 函数的 backward 方法创建辅助装饰器。

Autograd 函数是 torch.autograd.Function 的子类。确保 backward 以与 forward 相同的 autocast 状态执行。详见示例页面。

参数

device_type (str) – 要使用的设备类型。‘cuda’、‘cpu’、‘mtia’、‘xpu’ 等等。类型与 torch.device 的 type 属性相同。因此，您可以使用 Tensor.device.type 获取张量的设备类型。

class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)

参见 torch.autocast。

torch.cuda.amp.autocast(args...) 已弃用。请改用 torch.amp.autocast("cuda", args...)。

torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None)

torch.cuda.amp.custom_fwd(args...) 已弃用。请改用 torch.amp.custom_fwd(args..., device_type='cuda')。

torch.cuda.amp.custom_bwd(bwd)

torch.cuda.amp.custom_bwd(args...) 已弃用。请改用 torch.amp.custom_bwd(args..., device_type='cuda')。

class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True)

参见 torch.autocast。torch.cpu.amp.autocast(args...) 已弃用。请改用 torch.amp.autocast("cpu", args...)。

梯度缩放

如果某个算子的正向传播使用 float16 输入，则该算子的反向传播将产生 float16 梯度。幅度较小的梯度值可能无法在 float16 中表示。这些值将刷新为零（“下溢”），因此相应参数的更新将丢失。

为了防止下溢，“梯度缩放”将网络的损失乘以一个缩放因子，并在缩放后的损失上调用反向传播。通过网络反向流动的梯度随后会按相同因子进行缩放。换句话说，梯度值具有更大的幅度，因此它们不会刷新为零。

在优化器更新参数之前，每个参数的梯度（.grad 属性）应该被取消缩放，这样缩放因子就不会干扰学习率。

注意

AMP/fp16 可能不适用于每个模型！例如，大多数经过 bf16 预训练的模型无法在 fp16 最大值 65504 的数值范围内运行，会导致梯度溢出而不是下溢。在这种情况下，缩放因子可能会减小到小于 1，以尝试将梯度调整到 fp16 动态范围内可表示的数值。虽然人们可能期望缩放始终大于 1，但我们的 GradScaler 为了保持性能不提供此保证。如果您在使用 AMP/fp16 运行时在损失或梯度中遇到 NaNs，请验证您的模型是否兼容。

class torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)

参见 torch.amp.GradScaler。 torch.cuda.amp.GradScaler(args...) 已弃用。请改用 torch.amp.GradScaler("cuda", args...)。

class torch.cpu.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)

参见 torch.amp.GradScaler。torch.cpu.amp.GradScaler(args...) 已弃用。请改用 torch.amp.GradScaler("cpu", args...)。

Autocast Op 参考

Op 资格

以 float64 或非浮点 dtype 运行的 op 不符合资格，无论是否启用 autocast，它们都将以这些类型运行。

只有 out-of-place op 和 Tensor 方法符合资格。in-place 变体和显式提供 out=... Tensor 的调用在 autocast 启用区域中是允许的，但不会经过 autocasting。例如，在 autocast 启用区域中，a.addmm(b, c) 可以 autocast，但 a.addmm_(b, c) 和 a.addmm(b, c, out=d) 则不能。为了获得最佳性能和稳定性，在 autocast 启用区域中优先使用 out-of-place op。

使用显式 dtype=... 参数调用的 op 不符合资格，并将生成符合 dtype 参数的输出。

CUDA Op 特有行为

以下列表描述了 autocast 启用区域中符合资格的 op 的行为。无论这些 op 是作为 torch.nn.Module 的一部分、作为函数还是作为 torch.Tensor 方法调用，它们总是经过 autocasting。如果函数暴露在多个命名空间中，无论命名空间如何，它们都将经过 autocasting。

未在下面列出的 op 不会经过 autocasting。它们按照其输入定义的类型运行。但是，如果它们是 autocast op 的下游，autocasting 仍然可能改变未列出 op 的运行类型。

如果一个 op 未列出，我们假定它在 float16 中是数值稳定的。如果您认为未列出的 op 在 float16 中数值不稳定，请提交 issue。

可以 Autocast 到 float16 的 CUDA Ops

__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

可以 Autocast 到 float32 的 CUDA Ops

__pow__, __rdiv__, __rpow__, __rtruediv__, acos, asin, binary_cross_entropy_with_logits, cosh, cosine_embedding_loss, cdist, cosine_similarity, cross_entropy, cumprod, cumsum, dist, erfinv, exp, expm1, group_norm, hinge_embedding_loss, kl_div, l1_loss, layer_norm, log, log_softmax, log10, log1p, log2, margin_ranking_loss, mse_loss, multilabel_margin_loss, multi_margin_loss, nll_loss, norm, normalize, pdist, poisson_nll_loss, pow, prod, reciprocal, rsqrt, sinh, smooth_l1_loss, soft_margin_loss, softmax, softmin, softplus, sum, renorm, tan, triplet_margin_loss

Promote到最宽输入类型的 CUDA Ops

这些 op 不需要特定的 dtype 来保证稳定性，但它们接受多个输入并要求输入的 dtypes 匹配。如果所有输入都是 float16，则 op 在 float16 中运行。如果任何输入是 float32，autocast 会将所有输入转换为 float32 并使 op 在 float32 中运行。

addcdiv, addcmul, atan2, bilinear, cross, dot, grid_sample, index_put, scatter_add, tensordot

此处未列出的一些 op（例如，像 add 这样的二元 op）会原生地 promote 输入，而无需 autocasting 的干预。如果输入是 float16 和 float32 的混合，这些 op 会在 float32 中运行并生成 float32 输出，无论是否启用 autocast。

优先使用 binary_cross_entropy_with_logits 而非 binary_cross_entropy

torch.nn.functional.binary_cross_entropy()（以及包装它的 torch.nn.BCELoss）的 backward pass 可能会产生无法在 float16 中表示的梯度。在 autocast 启用区域中，forward 输入可能是 float16，这意味着 backward 梯度必须能在 float16 中表示（autocasting 将 float16 forward 输入转换为 float32 没有帮助，因为该转换在 backward 中必须被撤销）。因此，binary_cross_entropy 和 BCELoss 在 autocast 启用区域中会引发错误。

许多模型在二元交叉熵层之前使用 sigmoid 层。在这种情况下，请使用 torch.nn.functional.binary_cross_entropy_with_logits() 或 torch.nn.BCEWithLogitsLoss 结合这两个层。binary_cross_entropy_with_logits 和 BCEWithLogits 可以安全地进行 autocast。

XPU Op 特有行为 (实验性)

以下列表描述了 autocast 启用区域中符合资格的 op 的行为。无论这些 op 是作为 torch.nn.Module 的一部分、作为函数还是作为 torch.Tensor 方法调用，它们总是经过 autocasting。如果函数暴露在多个命名空间中，无论命名空间如何，它们都将经过 autocasting。

未在下面列出的 op 不会经过 autocasting。它们按照其输入定义的类型运行。但是，如果它们是 autocast op 的下游，autocasting 仍然可能改变未列出 op 的运行类型。

如果一个 op 未列出，我们假定它在 float16 中是数值稳定的。如果您认为未列出的 op 在 float16 中数值不稳定，请提交 issue。

可以 Autocast 到 float16 的 XPU Ops

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, RNNCell

可以 Autocast 到 float32 的 XPU Ops

__pow__, __rdiv__, __rpow__, __rtruediv__, binary_cross_entropy_with_logits, cosine_embedding_loss, cosine_similarity, cumsum, dist, exp, group_norm, hinge_embedding_loss, kl_div, l1_loss, layer_norm, log, log_softmax, margin_ranking_loss, nll_loss, normalize, poisson_nll_loss, pow, reciprocal, rsqrt, soft_margin_loss, softmax, softmin, sum, triplet_margin_loss

Promote到最宽输入类型的 XPU Ops

这些 op 不需要特定的 dtype 来保证稳定性，但它们接受多个输入并要求输入的 dtypes 匹配。如果所有输入都是 float16，则 op 在 float16 中运行。如果任何输入是 float32，autocast 会将所有输入转换为 float32 并使 op 在 float32 中运行。

bilinear, cross, grid_sample, index_put, scatter_add, tensordot

此处未列出的一些 op（例如，像 add 这样的二元 op）会原生地 promote 输入，而无需 autocasting 的干预。如果输入是 float16 和 float32 的混合，这些 op 会在 float32 中运行并生成 float32 输出，无论是否启用 autocast。

CPU Op 特有行为

以下列表描述了 autocast 启用区域中符合资格的 op 的行为。无论这些 op 是作为 torch.nn.Module 的一部分、作为函数还是作为 torch.Tensor 方法调用，它们总是经过 autocasting。如果函数暴露在多个命名空间中，无论命名空间如何，它们都将经过 autocasting。

未在下面列出的 op 不会经过 autocasting。它们按照其输入定义的类型运行。但是，如果它们是 autocast op 的下游，autocasting 仍然可能改变未列出 op 的运行类型。

如果一个 op 未列出，我们假定它在 bfloat16 中是数值稳定的。如果您认为未列出的 op 在 bfloat16 中数值不稳定，请提交 issue。float16 与 bfloat16 共享列表。

可以 Autocast 到 bfloat16 的 CPU Ops

conv1d, conv2d, conv3d, bmm, mm, linalg_vecdot, baddbmm, addmm, addbmm, linear, matmul, _convolution, conv_tbc, mkldnn_rnn_layer, conv_transpose1d, conv_transpose2d, conv_transpose3d, prelu, scaled_dot_product_attention, _native_multi_head_attention

可以 Autocast 到 float32 的 CPU Ops

avg_pool3d, binary_cross_entropy, grid_sampler, grid_sampler_2d, _grid_sampler_2d_cpu_fallback, grid_sampler_3d, polar, prod, quantile, nanquantile, stft, cdist, trace, view_as_complex, cholesky, cholesky_inverse, cholesky_solve, inverse, lu_solve, orgqr, inverse, ormqr, pinverse, max_pool3d, max_unpool2d, max_unpool3d, adaptive_avg_pool3d, reflection_pad1d, reflection_pad2d, replication_pad1d, replication_pad2d, replication_pad3d, mse_loss, cosine_embedding_loss, nll_loss, nll_loss2d, hinge_embedding_loss, poisson_nll_loss, cross_entropy_loss, l1_loss, huber_loss, margin_ranking_loss, soft_margin_loss, triplet_margin_loss, multi_margin_loss, ctc_loss, kl_div, multilabel_margin_loss, binary_cross_entropy_with_logits, fft_fft, fft_ifft, fft_fft2, fft_ifft2, fft_fftn, fft_ifftn, fft_rfft, fft_irfft, fft_rfft2, fft_irfft2, fft_rfftn, fft_irfftn, fft_hfft, fft_ihfft, linalg_cond, linalg_matrix_rank, linalg_solve, linalg_cholesky, linalg_svdvals, linalg_eigvals, linalg_eigvalsh, linalg_inv, linalg_householder_product, linalg_tensorinv, linalg_tensorsolve, fake_quantize_per_tensor_affine, geqrf, _lu_with_info, qr, svd, triangular_solve, fractional_max_pool2d, fractional_max_pool3d, adaptive_max_pool3d, multilabel_margin_loss_forward, linalg_qr, linalg_cholesky_ex, linalg_svd, linalg_eig, linalg_eigh, linalg_lstsq, linalg_inv_ex

Promote到最宽输入类型的 CPU Ops

这些 op 不需要特定的 dtype 来保证稳定性，但它们接受多个输入并要求输入的 dtypes 匹配。如果所有输入都是 bfloat16，则 op 在 bfloat16 中运行。如果任何输入是 float32，autocast 会将所有输入转换为 float32 并使 op 在 float32 中运行。

cat, stack, index_copy

此处未列出的一些 op（例如，像 add 这样的二元 op）会原生地 promote 输入，而无需 autocasting 的干预。如果输入是 bfloat16 和 float32 的混合，这些 op 会在 float32 中运行并生成 float32 输出，无论是否启用 autocast。