自动混合精度包 - torch.amp¶
torch.amp 提供了混合精度的便捷方法，
其中一些操作使用 torch.float32 (float) 数据类型，而其他操作
使用较低精度的浮点数据类型 (lower_precision_fp)：torch.float16 (half) 或 torch.bfloat16。一些操作，如线性层和卷积，
在 lower_precision_fp 中要快得多。其他操作，如归约，通常需要 float32 的动态范围。混合精度尝试将每个操作与其适当的数据类型匹配。
通常情况下，使用数据类型为 torch.float16 的“自动混合精度训练”会结合使用 torch.autocast 和
torch.cpu.amp.GradScaler 或 torch.cuda.amp.GradScaler，如 CUDA 自动混合精度示例
和 CUDA 自动混合精度配方 所示。
然而，torch.autocast 和 torch.GradScaler 是模块化的，如果需要，可以单独使用。
如 torch.autocast 的 CPU 示例部分所示，在 CPU 上使用数据类型为 torch.bfloat16 的“自动混合精度训练/推理”仅使用 torch.autocast。
对于CUDA和CPU，也分别提供了API：
torch.autocast("cuda", args...)等同于torch.cuda.amp.autocast(args...)。torch.autocast("cpu", args...)等同于torch.cpu.amp.autocast(args...)。对于 CPU，目前仅支持低精度的浮点数据类型torch.bfloat16。torch.GradScaler("cuda", args...)等同于torch.cuda.amp.GradScaler(args...)。torch.GradScaler("cpu", args...)等同于torch.cpu.amp.GradScaler(args...)。
torch.autocast 和 torch.cpu.amp.autocast 是在版本 1.10 中新增的。
自动转换¶
- class torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None)[源代码]¶
实例
autocast作为上下文管理器或装饰器，允许脚本中的区域以混合精度运行。在这些区域中，操作在由autocast选择的特定于操作的数据类型中运行，以提高性能同时保持精度。有关详细信息，请参阅Autocast Op Reference。
当进入启用自动转换的区域时，张量可以是任何类型。 在使用自动转换时，您不应在模型或输入上调用
half()或bfloat16()。autocast应该只包装你的网络的前向传递（包括损失计算）。不推荐在autocast下进行反向传递。反向操作运行在与autocast用于相应前向操作的相同类型中。CUDA 设备示例：
# 创建模型和优化器，使用默认精度 model = Net().cuda() optimizer = optim.SGD(model.parameters(), ...) for input, target in data: optimizer.zero_grad() # 为前向传播（模型 + 损失）启用自动类型转换 with torch.autocast(device_type="cuda"): output = model(input) loss = loss_fn(output, target) # 在调用backward()之前退出上下文管理器 loss.backward() optimizer.step()
参见 CUDA 自动混合精度示例 以了解在更复杂场景中的使用方法（例如，梯度惩罚、多个模型/损失、自定义 autograd 函数）。
autocast也可以用作装饰器，例如，在模型的forward方法上：class AutocastModel(nn.Module): ... @torch.autocast(device_type="cuda") def forward(self, input): ...
在启用自动转换的区域内生成的浮点张量可能是
float16。 返回禁用自动转换的区域后，将它们与不同数据类型的浮点张量一起使用可能会导致类型不匹配错误。如果发生这种情况，请将自动转换区域内生成的张量转换回float32（或根据需要转换为其他数据类型）。 如果来自自动转换区域的张量已经是float32，则转换操作不会执行任何操作，并且不会产生额外的开销。 CUDA 示例：# 在默认数据类型（这里假设为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 在 autocast 的操作列表中，应该以 float16 运行。 # 输入是 float32，但操作以 float16 运行并产生 float16 输出。 # 不需要手动转换。 e_float16 = torch.mm(a_float32, b_float32) # 也处理混合输入类型 f_float16 = torch.mm(d_float32, e_float16) # 退出 autocast 后，调用 f_float16.float() 以与 d_float32 一起使用 g_float32 = torch.mm(d_float32, f_float16.float())
CPU训练示例：
# 创建模型和优化器，使用默认精度 model = Net() optimizer = optim.SGD(model.parameters(), ...) for epoch in epochs: for input, target in data: optimizer.zero_grad() # 使用自动混合精度运行前向传播 with torch.autocast(device_type="cpu", dtype=torch.bfloat16): output = model(input) loss = loss_fn(output, target) loss.backward() optimizer.step()
CPU推理示例：
# 在默认精度下创建模型 model = Net().eval() with torch.autocast(device_type="cpu", dtype=torch.bfloat16): for input in data: # 使用自动混合精度运行前向传播。 output = model(input)
使用Jit Trace进行CPU推理示例：
```python 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() # 目前，我们建议禁用Jit Autocast Pass， # 由于问题：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) # 模型运行 for _ in range(3): model(torch.randn(1, input_size)) ```类型不匹配错误在启用自动转换的区域内是一个错误；如果你遇到这种情况，请提交一个问题。
autocast(enabled=False)子区域可以嵌套在启用了自动转换的区域中。 在本地禁用自动转换可能很有用，例如，如果您想强制子区域以特定的dtype运行。禁用自动转换可以让您对执行类型进行显式控制。在子区域中，来自周围区域的输入在使用前应转换为dtype：# 在默认数据类型（这里假设为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): # 调用 e_float16.float() 以确保 float32 执行 # （必要，因为 e_float16 是在自动转换区域内创建的） f_float32 = torch.mm(c_float32, e_float16.float()) # 重新进入自动转换启用区域时，不需要手动转换。 # torch.mm 再次以 float16 运行并产生 float16 输出，无论输入类型如何。 g_float16 = torch.mm(d_float32, f_float32)
自动转换状态是线程本地的。如果你想在新线程中启用它，必须在那个线程中调用上下文管理器或装饰器。这会影响
torch.nn.DataParallel和torch.nn.parallel.DistributedDataParallel当每个进程使用多个GPU时 （参见使用多个GPU）。- Parameters
device_type (str, 必需) – 要使用的设备类型。可能的值有：‘cuda’、‘cpu’、‘xpu’ 和 ‘hpu’。 该类型与
torch.device的 type 属性相同。 因此，您可以使用 Tensor.device.type 获取张量的设备类型。enabled (bool, 可选) – 是否在该区域启用自动转换。 默认值:
Truedtype (torch_dtype, 可选) – 是否使用 torch.float16 或 torch.bfloat16。
cache_enabled (bool, 可选) – 是否启用autocast内部的权重缓存。 默认值:
True
- class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)[源代码]¶
参见
torch.autocast。torch.cuda.amp.autocast(args...)等同于torch.autocast("cuda", args...)
- torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None)[源代码]¶
为自定义自动求导函数的
forward方法创建一个辅助装饰器。Autograd 函数是
torch.autograd.Function的子类。 请参阅 示例页面 以获取更多详细信息。- Parameters
cast_inputs (
torch.dtype或 None, 可选, 默认=None) – 如果不是None, 当forward在启用 autocast 的区域内运行时, 将传入的浮点型 CUDA Tensors 转换为目标 dtype (非浮点型 Tensors 不受影响), 然后禁用 autocast 执行forward。 如果为None,forward的内部操作将根据当前的 autocast 状态执行。
注意
如果装饰的
forward在自动转换区域外被调用，custom_fwd是一个无操作，并且cast_inputs没有效果。
- torch.cuda.amp.custom_bwd(bwd)[源代码]¶
为自定义自动求导函数的反向方法创建一个辅助装饰器。
Autograd 函数是
torch.autograd.Function的子类。 确保backward在执行时与forward具有相同的 autocast 状态。 有关更多详细信息，请参阅 示例页面。
- class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True)[源代码]¶
参见
torch.autocast.torch.cpu.amp.autocast(args...)等同于torch.autocast("cpu", args...)
梯度缩放¶
如果某个操作的前向传播输入为float16，则该操作的反向传播将产生float16梯度。
梯度值的幅度较小可能无法在float16中表示。
这些值将刷新为零（“下溢”），因此相应参数的更新将丢失。
为了防止下溢，“梯度缩放”将网络的损失乘以一个缩放因子，并对缩放后的损失进行反向传播。然后，通过网络反向流动的梯度也会按相同的因子进行缩放。换句话说，梯度值具有更大的幅度，因此它们不会变为零。
每个参数的梯度（.grad 属性）在优化器更新参数之前应该先进行缩放，以使缩放因子不会干扰学习率。
注意
AMP/fp16 可能不适用于所有模型！例如，大多数 bf16 预训练模型无法在 fp16 的最大数值范围 65504 内运行，并且会导致梯度溢出而不是下溢。在这种情况下，缩放因子可能会降至 1 以下，以尝试将梯度调整到 fp16 动态范围内可表示的数值。虽然人们可能期望缩放因子始终大于 1，但我们的 GradScaler 并不保证这一点以保持性能。如果在使用 AMP/fp16 运行时遇到损失或梯度中的 NaN，请验证您的模型是否兼容。
自动转换操作参考¶
操作资格¶
在 float64 或非浮点数数据类型中运行的操作不符合条件，并且无论是否启用自动转换，这些操作都会在这些类型中运行。
只有不在原位的操作和Tensor方法才符合条件。
在启用了autocast的区域内，允许使用就地变体和显式提供out=... Tensor的调用，但不会通过autocasting。
例如，在启用了autocast的区域内，a.addmm(b, c)可以进行autocast，
但a.addmm_(b, c)和a.addmm(b, c, out=d)不能。
为了获得最佳性能和稳定性，建议在启用了autocast的区域内使用不在原位的操作。
使用显式 dtype=... 参数调用的操作不符合条件，
并且将生成尊重 dtype 参数的输出。
CUDA 操作特定行为¶
以下列表描述了在启用自动转换区域中符合条件的操作的行为。
这些操作无论它们是作为torch.nn.Module的一部分调用，
还是作为函数调用，或是作为torch.Tensor方法调用，都会进行自动转换。如果函数在多个命名空间中暴露，
它们无论如何都会进行自动转换。
未列出的操作不会通过自动转换。它们以输入定义的类型运行。然而，如果未列出的操作位于自动转换操作的下游，自动转换仍可能改变这些操作运行的类型。
如果一个操作未列出，我们假设它在float16中是数值稳定的。
如果你认为一个未列出的操作在float16中是数值不稳定的，
请提交一个问题。
可以自动转换为float16的CUDA操作¶
__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操作¶
__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操作会提升到最宽的输入类型¶
这些操作不需要特定的dtype来保证稳定性，但需要多个输入，并要求输入的dtype匹配。如果所有输入都是float16，则操作以float16运行。如果任何输入是float32，autocast会将所有输入转换为float32并以float32运行操作。
addcdiv,
addcmul,
atan2,
bilinear,
cross,
dot,
grid_sample,
index_put,
scatter_add,
tensordot
此处未列出的一些操作（例如，像 add 这样的二元操作）在不需要自动转换的情况下，原生地提升输入。如果输入是 float16 和 float32 的混合，这些操作将以 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
可以安全地进行autocast。
CPU 操作特定行为¶
以下列表描述了在启用自动转换区域中符合条件的操作的行为。
这些操作无论它们是作为torch.nn.Module的一部分调用，
还是作为函数调用，或者是作为torch.Tensor方法调用，都会进行自动转换。如果函数在多个命名空间中暴露，
它们无论如何都会进行自动转换。
未列出的操作不会通过自动转换。它们以输入定义的类型运行。然而，如果未列出的操作位于自动转换操作的下游，自动转换仍可能改变这些操作运行的类型。
如果一个操作未列出，我们假设它在bfloat16中是数值稳定的。
如果你认为一个未列出的操作在bfloat16中是数值不稳定的，
请提交一个问题。
可以自动转换为bfloat16的CPU操作¶
conv1d,
conv2d,
conv3d,
bmm,
mm,
baddbmm,
addmm,
addbmm,
linear,
matmul,
_convolution
可以自动转换为float32的CPU操作¶
conv_transpose1d,
conv_transpose2d,
conv_transpose3d,
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,
ctc_loss,
kl_div,
multilabel_margin_loss,
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_matrix_norm,
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,
eig,
geqrf,
lstsq,
_lu_with_info,
qr,
solve,
svd,
symeig,
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操作¶
这些操作不需要特定的数据类型以确保稳定性，但需要多个输入，并要求输入的数据类型匹配。如果所有输入都是bfloat16，则操作将以bfloat16运行。如果任何输入是float32，自动转换会将所有输入转换为float32，并以float32运行操作。
cat,
stack,
index_copy
此处未列出的一些操作（例如，像 add 这样的二元操作）在无需自动转换的情况下，原生地提升输入。如果输入是 bfloat16 和 float32 的混合，这些操作将以 float32 运行并生成 float32 输出，无论是否启用了自动转换。
