自动混合精度包 - 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的实例用作上下文管理器或装饰器，允许脚本中的某些区域以混合精度运行。在这些区域中，ops 以自动转换选择的特定操作 dtype 运行， 以提高性能同时保持准确性。 有关详细信息，请参阅 自动转换操作参考。
在进入自动混合精度启用区域时，张量可以是任何类型。 在使用自动类型转换时，不应在模型或输入上调用
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 自动混合精度示例以了解在更复杂场景中的使用（包括梯度缩放）， 例如，梯度惩罚、多个模型/损失函数、自定义 autograd 函数。
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)
CPU 推理示例与 Jit 追踪：
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))
类型不匹配错误是在自动混合精度区域中的一个 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状态是线程本地的。如果你想在新线程中启用它，必须在该线程中调用上下文管理器或装饰器。这会影响
torch.nn.DataParallel和torch.nn.parallel.DistributedDataParallel，当与每个进程多个GPU一起使用时（参见使用多个GPU）。- Parameters
设备类型 (str, 必需) – 使用的设备类型。可能的值为：‘cuda’，‘cpu’，‘xpu’ 和 ‘hpu’。 该类型与type属性相同
torch.device。 因此，您可以使用Tensor.device.type获取张量的设备类型。启用 (bool, 可选) – 是否应在该区域启用自动类型转换。 默认值：
Truedtype (torch_dtype, 可选) – 是否使用 torch.float16 或 torch.bfloat16。
启用缓存 (bool, 可选) – 是否启用autocast中的权重缓存。 默认值：
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的子类。 详情请参见 示例页面。- Parameters
cast_inputs (
torch.dtype或 None，可选，默认值=None) – 如果不是None, 当forward在启用了自动类型转换的区域运行时，会将传入的浮点型 CUDA 张量转换为目标数据类型（非浮点型张量不受影响）, 然后在禁用自动类型转换的情况下执行forward。 如果为None，forward的内部操作将使用当前的自动类型转换状态执行。
注意
如果装饰的
forward在未启用autocast的区域中被调用，custom_fwd将不起作用，并且cast_inputs没有任何效果。
- torch.cuda.amp.custom_bwd(bwd)[source]¶
为自定义自动微分函数的反向方法创建一个辅助装饰器。
自动求导函数是
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预训练模型无法在最大值为65504的fp16数值范围内运行，并且会导致梯度溢出而不是下溢。在这种情况下，缩放因子可能会降低到小于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中数值不稳定，
请提交一个问题。
CUDA 运算符可以自动转换为 float16¶
__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
CUDA 运算符可以自动转换为 float32¶
__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，
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 并不解决问题，因为在反向传播中需要将其还原）。
因此，在自动混合精度启用的区域中，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中数值不稳定，
请提交一个问题。
可以在自动类型转换中使用的CPU运算符bfloat16¶
conv1d,
conv2d,
conv3d,
bmm,
mm,
baddbmm,
addmm,
addbmm,
linear,
matmul,
_convolution
可以在自动类型转换中使用的CPU运算符float32¶
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，
autocast 将把所有输入转换为 float32 并在 float32 中运行操作。
cat,
stack,
index_copy
有些操作在这里没有列出（例如，二元操作如add）会原生地提升输入，而无需自动类型转换的干预。如果输入是bfloat16和float32的混合，这些操作会在float32中运行，并产生float32输出，无论是否启用了自动类型转换。