自动混合精度包 - torch.amp

torch.amp 为混合精度提供便捷方法，其中某些操作使用 torch.float32 (float) 数据类型，而其他操作使用较低精度浮点数据类型 (lower_precision_fp)：torch.float16 (half) 或 torch.bfloat16。 某些操作（如线性层和卷积）在 lower_precision_fp 中速度更快。 其他操作（如归约）通常需要 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 上是否可用。

参数

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

返回类型

bool

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

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

在这些区域中，操作以自动类型转换选择的操作特定数据类型运行，以提高性能并保持准确性。 有关详细信息，请参阅 自动类型转换操作参考。

当进入启用自动类型转换的区域时，张量可以是任何类型。 使用自动类型转换时，你不应该在你的模型或输入上调用 half() 或 bfloat16()。

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()

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

autocast 也可以用作装饰器，例如，在你模型的前向方法上

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

在启用自动类型转换的区域中生成的浮点张量可能是 float16。 返回到禁用自动类型转换的区域后，将它们与不同数据类型的浮点张量一起使用可能会导致类型不匹配错误。 如果是这样，将自动类型转换区域中生成的张量转换回 float32（或者如果需要，转换为其他数据类型）。 如果来自自动类型转换区域的张量已经是 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 跟踪的 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(enabled=False) 子区域可以嵌套在启用自动类型转换的区域中。 本地禁用自动类型转换可能很有用，例如，如果您想强制子区域以特定的 dtype 运行。 禁用自动类型转换使您可以显式控制执行类型。 在子区域中，来自周围区域的输入应在使用前转换为 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)

自动类型转换状态是线程本地的。 如果您希望在新线程中启用它，则必须在该线程中调用上下文管理器或装饰器。 当每个进程使用多个 GPU 时，这会影响 torch.nn.DataParallel 和 torch.nn.parallel.DistributedDataParallel（请参阅使用多个 GPU）。

参数

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

• enabled (bool, 可选) – 是否应在该区域中启用自动类型转换。 默认值：True

• dtype (torch_dtype, 可选) – 在自动类型转换中运行的操作的数据类型。 如果 dtype 为 None，则它使用默认值（CUDA 为 torch.float16，CPU 为 torch.bfloat16），由 get_autocast_dtype() 给出。 默认值：None

• cache_enabled (bool, 可选) – 是否应启用自动类型转换内部的权重缓存。 默认值：True

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

为自定义自动求导函数的前向方法创建辅助装饰器。

自动求导函数是 torch.autograd.Function 的子类。 有关更多详细信息，请参阅示例页面。

参数

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

• cast_inputs (torch.dtype 或 None，可选，默认值=None) – 如果不是 None，则当前向传播在启用自动类型转换的区域中运行时，将传入的浮点张量转换为目标数据类型（非浮点张量不受影响），然后执行禁用自动类型转换的前向传播。 如果为 None，则 forward 的内部操作以当前的自动类型转换状态执行。

注意

如果修饰后的 forward 在启用自动类型转换的区域之外被调用，则 custom_fwd 是空操作，并且 cast_inputs 没有效果。

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

为自定义自动求导函数的反向方法创建辅助装饰器。

自动求导函数是 torch.autograd.Function 的子类。 确保 backward 以与 forward 相同的自动类型转换状态执行。 有关更多详细信息，请参阅示例页面。

参数

device_type (str) – 要使用的设备类型。 ‘cuda’、‘cpu’、‘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 预训练模型无法在最大值 65504 的 fp16 数值范围内运行，并且会导致梯度溢出而不是下溢。 在这种情况下，比例因子可能会降至 1 以下，以尝试将梯度带到 fp16 动态范围中可表示的数字。 虽然人们可能期望比例始终高于 1，但我们的 GradScaler 不会保证维持性能。 如果在使用 AMP/fp16 运行时遇到损失或梯度中出现 NaN，请验证您的模型是否兼容。

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...) 代替。

自动类型转换操作参考

Op 资格

在 float64 或非浮点 dtype 中运行的 Ops 不符合条件，无论是否启用自动类型转换，都将以这些类型运行。

只有异地 (out-of-place) 操作和 Tensor 方法符合条件。在启用自动类型转换的区域中，允许使用原地 (in-place) 变体和显式提供 out=... Tensor 的调用，但不会进行自动类型转换。例如，在启用自动类型转换的区域中，a.addmm(b, c) 可以自动类型转换，但 a.addmm_(b, c) 和 a.addmm(b, c, out=d) 则不能。为了获得最佳性能和稳定性，请在启用自动类型转换的区域中首选异地操作。

使用显式 dtype=... 参数调用的 Ops 不符合条件，并将生成尊重 dtype 参数的输出。

CUDA Op 特定行为

以下列表描述了在启用自动类型转换的区域中符合条件的 Ops 的行为。无论这些 Ops 是作为 torch.nn.Module 的一部分、作为函数还是作为 torch.Tensor 方法调用，它们始终会进行自动类型转换。如果函数在多个命名空间中公开，则无论命名空间如何，它们都会进行自动类型转换。

下面未列出的 Ops 不会进行自动类型转换。它们以其输入定义的类型运行。但是，如果未列出的 Ops 位于自动类型转换的 Ops 的下游，自动类型转换仍可能更改这些 Ops 运行的类型。

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

可以自动类型转换为 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

可以自动类型转换为 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

CUDA Ops，可以提升到最宽的输入类型

这些 Ops 对稳定性没有特定的 dtype 要求，但接受多个输入，并要求输入的 dtype 匹配。如果所有输入都是 float16，则 op 在 float16 中运行。如果任何输入是 float32，则自动类型转换会将所有输入转换为 float32，并在 float32 中运行 op。

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

此处未列出的一些 Ops (例如，像 add 这样的二元操作) 在没有自动类型转换干预的情况下，自然会提升输入类型。如果输入是 float16 和 float32 的混合，则无论是否启用自动类型转换，这些 ops 都会在 float32 中运行并生成 float32 输出。

优先选择 binary_cross_entropy_with_logits 而不是 binary_cross_entropy

torch.nn.functional.binary_cross_entropy() (以及包装它的 torch.nn.BCELoss) 的反向传播可能会产生无法在 float16 中表示的梯度。在启用自动类型转换的区域中，前向输入可能是 float16，这意味着反向梯度必须可以在 float16 中表示 (将 float16 前向输入自动类型转换为 float32 没有帮助，因为该转换必须在反向传播中逆转)。因此，binary_cross_entropy 和 BCELoss 在启用自动类型转换的区域中会引发错误。

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

XPU Op 特定行为 (实验性)

以下列表描述了在启用自动类型转换的区域中符合条件的 Ops 的行为。无论这些 Ops 是作为 torch.nn.Module 的一部分、作为函数还是作为 torch.Tensor 方法调用，它们始终会进行自动类型转换。如果函数在多个命名空间中公开，则无论命名空间如何，它们都会进行自动类型转换。

下面未列出的 Ops 不会进行自动类型转换。它们以其输入定义的类型运行。但是，如果未列出的 Ops 位于自动类型转换的 Ops 的下游，自动类型转换仍可能更改这些 Ops 运行的类型。

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

可以自动类型转换为 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

可以自动类型转换为 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

XPU Ops，可以提升到最宽的输入类型

这些 Ops 对稳定性没有特定的 dtype 要求，但接受多个输入，并要求输入的 dtype 匹配。如果所有输入都是 float16，则 op 在 float16 中运行。如果任何输入是 float32，则自动类型转换会将所有输入转换为 float32，并在 float32 中运行 op。

bilinear, cross, grid_sample, index_put, scatter_add, tensordot

此处未列出的一些 Ops (例如，像 add 这样的二元操作) 在没有自动类型转换干预的情况下，自然会提升输入类型。如果输入是 float16 和 float32 的混合，则无论是否启用自动类型转换，这些 ops 都会在 float32 中运行并生成 float32 输出。

CPU Op 特定行为

以下列表描述了在启用自动类型转换的区域中符合条件的 Ops 的行为。无论这些 Ops 是作为 torch.nn.Module 的一部分、作为函数还是作为 torch.Tensor 方法调用，它们始终会进行自动类型转换。如果函数在多个命名空间中公开，则无论命名空间如何，它们都会进行自动类型转换。

下面未列出的 Ops 不会进行自动类型转换。它们以其输入定义的类型运行。但是，如果未列出的 Ops 位于自动类型转换的 Ops 的下游，自动类型转换仍可能更改这些 Ops 运行的类型。

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

可以自动类型转换为 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

可以自动类型转换为 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

CPU Ops，可以提升到最宽的输入类型

这些 Ops 对稳定性没有特定的 dtype 要求，但接受多个输入，并要求输入的 dtype 匹配。如果所有输入都是 bfloat16，则 op 在 bfloat16 中运行。如果任何输入是 float32，则自动类型转换会将所有输入转换为 float32，并在 float32 中运行 op。

cat, stack, index_copy

此处未列出的一些 Ops (例如，像 add 这样的二元操作) 在没有自动类型转换干预的情况下，自然会提升输入类型。如果输入是 bfloat16 和 float32 的混合，则无论是否启用自动类型转换，这些 ops 都会在 float32 中运行并生成 float32 输出。