自动混合精度包 - 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)[source]¶
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()
有关在更复杂场景（例如梯度惩罚、多个模型/损失、自定义自动梯度函数）中使用（以及梯度缩放）的示例，请参阅 CUDA 自动混合精度示例。
autocast也可以用作装饰器，例如，在模型的forward方法上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)
自动转换状态是线程局部的。如果您想在新线程中启用它，则必须在该线程中调用上下文管理器或装饰器。这会影响
torch.nn.DataParallel和torch.nn.parallel.DistributedDataParallel在每个进程使用多个 GPU 时（请参阅 使用多个 GPU）。- 参数
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, 可选) – 是否启用自动类型转换内部的权重缓存。默认值：
True
- class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)[source]¶
参见
torch.autocast.torch.cuda.amp.autocast(args...)等效于torch.autocast("cuda", args...)
- torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None)[source]¶
为自定义自动微分函数的
forward方法创建一个辅助装饰器。自动微分函数是
torch.autograd.Function的子类。有关更多详细信息，请参阅 示例页面。- 参数
cast_inputs (
torch.dtype或 None，可选，默认值=None) – 如果不为None，当forward在启用自动广播的区域中运行时，将传入的浮点 CUDA 张量转换为目标数据类型（非浮点张量不受影响），然后在禁用自动广播的情况下执行forward。如果为None，则forward的内部操作将使用当前自动广播状态执行。
注意
如果装饰的
forward在启用自动广播的区域之外调用，则custom_fwd为无操作，并且cast_inputs不起作用。
- torch.cuda.amp.custom_bwd(bwd)[source]¶
为自定义自动微分函数的 backward 方法创建辅助装饰器。
自动微分函数是
torch.autograd.Function的子类。确保backward以与forward相同的自动转换状态执行。有关更多详细信息，请参见 示例页面。
- class torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True)[source]¶
参见
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 或非浮点类型中运行的操作不符合条件，无论是否启用自动转换，它们都将以这些类型运行。
只有非就地操作和张量方法符合条件。就地变体和显式提供 out=... 张量的调用在启用自动转换的区域中是允许的，但不会经过自动转换。例如，在启用自动转换的区域中，a.addmm(b, c) 可以自动转换，但 a.addmm_(b, c) 和 a.addmm(b, c, out=d) 不能。为了获得最佳性能和稳定性，在启用自动转换的区域中优先使用非就地操作。
使用显式 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 操作¶
这些操作不需要特定的数据类型来保证稳定性，但它们接受多个输入，并要求输入的数据类型匹配。如果所有输入都是 float16，则操作在 float16 中运行。如果任何输入是 float32，则自动转换将所有输入转换为 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 对自动转换是安全的。
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 输出，无论自动转换是否启用。
