自动混合精度包 - 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)[source][source] 说明
返回一个布尔值，指示是否在
device_type上可用自动类型转换。- 参数:
device_type (str) – 要使用的设备类型。可能的值有：‘cuda’，‘cpu’，‘mtia’，‘xpu’等。类型与
torch.device的类型属性相同。因此，您可以使用 Tensor.device.type 获取张量的设备类型。- 返回类型:
- class torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None)[source][source]¶
autocast的实例用作上下文管理器或装饰器，允许脚本中的某些区域以混合精度运行。在这些区域中，操作以 autocast 选择的特定操作 dtype 运行，以提高性能同时保持精度。有关详细信息，请参阅 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()
请参阅自动混合精度示例，了解在更复杂场景（例如梯度惩罚、多个模型/损失、自定义 autograd 函数）中的使用（包括梯度缩放）。
也可以用作装饰器，例如用于模型的
forward方法：class AutocastModel(nn.Module): ... @torch.autocast(device_type="cuda") def forward(self, input): ...
在自动转换区域产生的浮点张量可能为
float16。返回到自动转换禁用区域后，使用它们与不同 dtypes 的浮点张量可能会导致类型不匹配错误。如果是这样，请将自动转换区域产生的张量转换回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 Trace
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)子区域可以嵌套在自动类型转换启用的区域内。在需要强制子区域在特定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, required) – 要使用的设备类型。可能的值有：‘cuda’，‘cpu’，‘mtia’，‘xpu’，和‘hpu’。类型与
torch.device的类型属性相同。因此，您可以使用 Tensor.device.type 获取张量的设备类型。enabled (bool, optional) – 是否在该区域启用自动转换。默认：
Truedtype (torch_dtype, optional) – 自动转换中运行的 ops 的数据类型。如果
dtype是None，则使用默认值（CUDA 为torch.float16，CPU 为torch.bfloat16），由get_autocast_dtype()给出。默认：Nonecache_enabled (bool, 可选) – 是否启用 autocast 内部的权重缓存。默认：
True
- torch.amp.custom_fwd(fwd=None, *, device_type, cast_inputs=None)[source][source]¶
创建一个用于自定义 autograd 函数的
forward方法的辅助装饰器。Autograd 函数是
torch.autograd.Function的子类。请参阅示例页面以获取更多详细信息。- 参数:
device_type (str) – 要使用的设备类型。‘cuda’，‘cpu’，‘mtia’，‘xpu’等。类型与
torch.device的 type 属性相同。因此，您可以使用 Tensor.device.type 获取张量的设备类型。cast_inputs (
torch.dtype或 None, 可选，默认=None) – 如果不是None，当forward在启用自动转换的区域运行时，将传入的浮点张量转换为目标数据类型（非浮点张量不受影响），然后禁用自动转换执行forward。如果None，forward的内部操作以当前自动转换状态执行。
备注
如果装饰的
forward在启用自动转换的区域外被调用，custom_fwd是一个空操作，cast_inputs没有效果。
- torch.amp.custom_bwd(bwd=None, *, device_type)[source][source]¶
创建自定义自动微分函数反向方法的辅助装饰器。
自动微分函数是
torch.autograd.Function的子类。确保backward在与forward相同的 autocast 状态下执行。请参阅示例页面以获取更多详细信息。- 参数:
device_type (str) – 要使用的设备类型。‘cuda’，‘cpu’，‘mtia’，‘xpu’等。类型与
torch.device的类型属性相同。因此，您可以使用 Tensor.device.type 获取张量的设备类型。
- class torch.cuda.amp.autocast(enabled=True, dtype=torch.float16, cache_enabled=True)[source][source]¶
见
torch.autocast。torch.cuda.amp.autocast(args...)已弃用。请使用torch.amp.autocast("cuda", args...)代替。
- torch.cuda.amp.custom_fwd(fwd=None, *, cast_inputs=None)[source][source]¶
torch.cuda.amp.custom_fwd(args...)已弃用。请使用torch.amp.custom_fwd(args..., device_type='cuda')代替。
- torch.cuda.amp.custom_bwd(bwd)[source][source]¶
torch.cuda.amp.custom_bwd(args...)已弃用。请使用torch.amp.custom_bwd(args..., device_type='cuda')代替。
- class torch.cpu.amp.autocast(启用=True, 数据类型=torch.bfloat16, 缓存启用=True)[source][source] ¶
请参阅
torch.autocast。torch.cpu.amp.autocast(args...)已弃用。请使用torch.amp.autocast("cpu", args...)代替。
梯度缩放
如果特定操作的 forward 过程有 float16 个输入，则该操作的 backward 过程将产生 float16 个梯度。梯度值可能无法在 float16 中表示，其值将清零（下溢），因此相应参数的更新将丢失。
为了防止下溢，“梯度缩放”会将网络的损失乘以一个缩放因子，并对缩放后的损失进行反向传播。然后，通过网络反向流动的梯度也会按相同的因子进行缩放。换句话说，梯度值具有更大的幅度，因此不会冲刷到零。
在优化器更新参数之前，每个参数的梯度（ .grad 属性）应该进行未缩放处理，以确保缩放因子不会干扰学习率。
备注
AMP/fp16 可能不适用于每个模型！例如，大多数 bf16 预训练模型无法在 fp16 数值范围内运行，该范围最大为 65504，这会导致梯度上溢而不是下溢。在这种情况下，缩放因子可能会小于 1，作为将梯度转换为 fp16 动态范围内可表示的数字的尝试。虽然人们可能期望缩放因子始终大于 1，但我们的 GradScaler 并不保证这一点以保持性能。如果在运行 AMP/fp16 时遇到损失或梯度中的 NaN，请验证您的模型是否兼容。
自动转换操作参考 ¶
沉浸式翻译
在 float64 或非浮点数据类型上运行的运算符不符合条件，无论是否启用自动类型转换，都会以这些类型运行。
只有就地运算符和 Tensor 方法符合条件。在自动类型转换区域中，允许就地变体和显式提供 out=... Tensor 的调用，但不会进行自动类型转换。例如，在自动类型转换区域中， 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 可以安全地进行自动类型转换。
XPU 操作特定行为（实验性）¶
以下列表描述了在自动转换启用区域中符合条件的操作的行为。这些操作无论作为 torch.nn.Module 的一部分、作为一个函数还是作为一个 torch.Tensor 方法被调用，都会始终通过自动转换。如果函数在多个命名空间中公开，它们无论在哪个命名空间都会通过自动转换。
未列出的操作不会通过自动转换。它们将在由它们的输入定义的类型中运行。然而，如果未列出的操作在自动转换操作之后，自动转换仍可能改变它们运行的类型。
如果一个操作未列出，我们假设它在 float16 中是数值稳定的。如果您认为一个未列出的操作在 float16 中是数值不稳定的，请提交一个问题。
支持自动转换为 float16 的 XPU 操作
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 操作
__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 操作
这些操作不需要特定的 dtype 以保证稳定性，但需要多个输入，并且要求输入的 dtype 匹配。如果所有输入都是 float16 ，则操作在 float16 下运行。如果任何输入是 float32 ，则自动转换将所有输入转换为 float32 ，并在 float32 下运行操作。
bilinear,
cross,
grid_sample,
index_put,
scatter_add,
tensordot
一些未在此列出的操作（例如，二进制操作 add ）会原生提升输入，无需自动类型转换的干预。如果输入是 float16 和 float32 的混合，这些操作将在 float32 中运行并产生 float32 输出，无论自动类型转换是否启用。
CPU 操作特定行为
以下列表描述了在自动类型转换启用区域中符合条件的操作的行为。这些操作无论作为 torch.nn.Module 的一部分、作为函数还是作为 torch.Tensor 方法被调用，都会始终通过自动类型转换。如果函数在多个命名空间中公开，无论命名空间如何，都会通过自动类型转换。
未列出的操作不会通过自动类型转换。它们将在由其输入定义的类型中运行。然而，如果未列出的操作是自动类型转换操作的下级，自动类型转换仍可能更改它们运行的类型。
如果一个操作未列出，我们假设它在 bfloat16 中数值稳定。如果您认为一个未列出的操作在 bfloat16 中数值不稳定，请提交问题。 float16 共享了 bfloat16 的列表。
可以自动转换到 bfloat16 的 CPU 操作 bfloat16
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 操作 float32
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 操作@0#
这些操作不需要特定的数据类型以保证稳定性，但需要多个输入，并且要求输入的数据类型匹配。如果所有输入都是 bfloat16 ，则操作在 bfloat16 下运行。如果任何输入是 float32 ，则 autocast 将所有输入转换为 float32 ，并在 float32 下运行该操作。
cat,
stack,
index_copy
列表中未列出的某些操作（例如，二进制操作 add ）会原生提升输入，无需 autocast 干预。如果输入是 bfloat16 和 float32 的混合，则这些操作在 float32 下运行，并产生 float32 输出，无论 autocast 是否启用。
