自动混合精度软件包 - torch.amp¶
为混合精度提供便捷的方法，
其中，某些操作使用 （） 数据类型，而其他操作
使用低精度浮点数据类型 （）：（） 或 .一些运算，如线性层和卷积、
在 中要快得多。其他操作（如缩减）通常需要动态
的范围。混合精度尝试将每个运算与其适当的数据类型匹配。torch.float32floatlower_precision_fptorch.float16halftorch.bfloat16lower_precision_fpfloat32
通常，“自动混合精度训练”的数据类型为 uses 和 together，如 Automatic Mixed Precision 示例和 Automatic Mixed Precision 配方所示。
但是，并且是模块化的，如果需要，可以单独使用。
如 CPU 示例部分所示，在 CPU 上使用 “自动混合精度训练/推理”
数据类型仅使用 .torch.float16torch.amp.GradScalertorch.GradScalertorch.bfloat16
警告
torch.cuda.amp.autocast(args...)，将被弃用。请改用 or。 ，将被弃用。请改用 or。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 的 Torch 类autocast（device_type， dtype=None， enabled=True， cache_enabled=None）[来源]¶
的
实例充当上下文管理器或装饰器，这些 允许脚本的区域以混合精度运行。在这些区域中，运算在 autocast 选择的特定于运算的 dtype 中运行 提高性能，同时保持准确性。 有关详细信息，请参阅 Autocast Op Reference （自动转换操作参考）。
当进入启用自动转换的区域时，Tensor 可以是任何类型。 使用自动转换时，不应对模型或输入调用 或。
half()bfloat16()应仅包装网络的 forward pass，包括 loss 计算。不建议在 autocast 下向后传递。 Backward 操作以与 autocast 用于相应 forward 操作的相同类型运行。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 函数）。
class AutocastModel(nn.Module): ... @torch.autocast(device_type="cuda") def forward(self, input): ...
在启用自动转换的区域中生成的浮点 Tensor 可能是 。 返回到禁用自动转换的区域后，将它们与 floating-point 一起使用 不同 dtype 的张量可能会导致类型不匹配错误。如果是这样，则强制转换 Tensor 在 autocast 区域中生成，返回到（或根据需要返回其他 dtype）。 如果来自自动转换区域的 Tensor 已经是 ，则强制转换是无操作， 并且不会产生额外的开销。 CUDA 示例：
float16float32float32# 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))
启用自动转换的区域中的类型不匹配错误是一个错误;如果这是你观察到的， 请提交 issue。
autocast(enabled=False)子区域可以嵌套在启用了自动转换的区域中。 在本地禁用自动转换可能很有用，例如，如果要强制使用子区域 在特定的 .禁用自动广播可以让你对 执行类型。在子区域中，来自周围区域的输入 应在使用前强制转换：dtypedtype# 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 状态是线程本地的。如果你想在新线程中启用它，上下文管理器或装饰器 必须在该线程中调用。这会影响
每个进程与多个 GPU 一起使用时 （请参阅使用多个 GPU）。- 参数
device_type （str， required） – 要使用的设备类型。可能的值为：'cuda'、'cpu'、'xpu' 和 'hpu'。 type 与 的 type 属性相同
。 因此，您可以使用 Tensor.device.type 获取 Tensor 的设备类型。enabled （bool， optional） – 是否应在区域中启用自动转换。 违约：
Truedtype （torch_dtype，可选） – 在自动转换中运行的操作的数据类型。它使用默认值 （对于 CUDA 和 CPU），由 ，如果
为 。 违约：torch.float16torch.bfloat16get_autocast_dtype()NoneNonecache_enabled （bool， optional） – 是否应启用自动转换中的权重缓存。 违约：
True
- torch.amp。custom_fwd（fwd=无， *， device_type， cast_inputs=无）[来源]¶
为自定义 autograd 函数的方法创建一个辅助装饰器。
forwardAutograd 函数是 的子类
。 有关更多详细信息，请参阅示例页面。- 参数
device_type （str） – 要使用的设备类型。'cuda'、'cpu'、'xpu' 等。 type 与 的 type 属性相同
。 因此，您可以使用 Tensor.device.type 获取 Tensor 的设备类型。cast_inputs （
或 None， optional， default=None） – 如果不是 ， 在启用了自动广播的区域运行时，将转换传入的 浮点 Tensor 添加到目标 dtype 中（非浮点 Tensor 不受影响）， 然后在禁用自动转换的情况下执行。 如果 ， 的内部操作以当前自动转换状态执行。NoneforwardforwardNoneforward
- torch.amp。custom_bwd（bwd=无， *， device_type）[来源]¶
为自定义 autograd 函数的向后方法创建一个辅助装饰器。
Autograd 函数是 的子类
。 确保 以与 相同的自动转换状态执行。 有关更多详细信息，请参阅示例页面。backwardforward- 参数
device_type （str） – 要使用的设备类型。'cuda'、'cpu'、'xpu' 等。 type 与 的 type 属性相同
。 因此，您可以使用 Tensor.device.type 获取 Tensor 的设备类型。
- 类 torch.cuda.amp。autocast（enabled=True， dtype=torch.float16， cache_enabled=True）[来源]¶
-
torch.cuda.amp.autocast(args...)已弃用。请改用。torch.amp.autocast("cuda", args...)
- torch.cuda.amp。custom_fwd（fwd=无， *， cast_inputs=无）[来源]¶
torch.cuda.amp.custom_fwd(args...)已弃用。请改用。torch.amp.custom_fwd(args..., device_type='cuda')
梯度缩放¶
如果特定运算的正向传递有输入，则
该运算将产生梯度。
具有较小量级的梯度值可能无法在 中表示。
这些值将刷新为零（“下溢”），因此相应参数的更新将丢失。float16float16float16
为了防止下溢，“梯度缩放”将网络的损失乘以比例因子，并且 对缩放的损失调用向后传递。通过网络向后流动的梯度是 然后按相同的因子进行缩放。换句话说，梯度值具有更大的幅度， 因此，它们不会刷新为零。
每个参数的 gradient （ attribute） 应在优化器之前取消缩放
更新参数，以便 Scale Factor 不会干扰学习率。.grad
注意
AMP/fp16 可能不适用于每个型号！例如，大多数 bf16 预训练模型无法在 FP16 数值范围最大为 65504，将导致梯度上溢而不是下溢。在 在这种情况下，比例因子可能会减小到 1 以下，以尝试将渐变变为数字 可在 FP16 动态范围内表示。虽然人们可能期望小数位数始终高于 1，但我们的 GradScaler 不保证保持性能。如果您在损失中遇到 NaN 或梯度（在使用 AMP/fp16 运行时）验证您的模型是否兼容。
Autocast Op 参考¶
Op 资格¶
在 dtype 中运行或非浮点 dtype 的运算不符合条件，并且将
无论是否启用了 Autocast，都以这些类型运行。float64
只有 out-of-place operations 和 Tensor 方法才符合条件。
显式提供 Tensor 的就地变体和调用
允许在启用自动转换的区域中使用，但不会进行自动转换。
例如，在启用了自动转换的区域可以自动转换，
但是不能。
为了获得最佳性能和稳定性，请在启用了 autocast 的情况下首选 off-of-place operations
地区。out=...a.addmm(b, c)a.addmm_(b, c)a.addmm(b, c, out=d)
使用显式参数调用的操作不符合条件，
并将生成尊重参数的输出。dtype=...dtype
CUDA Op 特定行为¶
以下列表描述了启用自动转换的区域中符合条件的操作的行为。
这些操作始终会进行自动转换，无论它们是作为 ，
作为函数或方法。如果函数在多个命名空间中公开，则
无论命名空间如何，它们都会进行自动转换。
下面未列出的 Ops 不会进行自动转换。它们在 由他们的输入定义。但是，自动转换可能仍会更改类型 如果未列出的操作位于自动转换操作的下游，则运行其中的未列出操作。
如果一个运算未列出，我们假设它在 中的数值稳定。
如果您认为未列出的操作在 中数值不稳定 ，
请提交 issue。float16float16
可以自动转换为
__matmul__, , , , , , , , , , , , , , , , , , , , , ,addbmmaddmmaddmvaddrbaddbmmbmmchain_matmulmulti_dotconv1dconv2dconv3dconv_transpose1dconv_transpose2dconv_transpose3dGRUCelllinearLSTMCellmatmulmmmvpreluRNNCell
可以自动转换为
__pow__, , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , ,__rdiv____rpow____rtruediv__acosasinbinary_cross_entropy_with_logitscoshcosine_embedding_losscdistcosine_similaritycross_entropycumprodcumsumdisterfinvexpexpm1group_normhinge_embedding_losskl_divl1_losslayer_normloglog_softmaxlog10log1plog2margin_ranking_lossmse_lossmultilabel_margin_lossmulti_margin_lossnll_lossnormnormalizepdistpoisson_nll_losspowprodreciprocalrsqrtsinhsmooth_l1_losssoft_margin_losssoftmaxsoftminsoftplussumrenormtantriplet_margin_loss
提升到最宽输入类型的 CUDA Ops¶
这些 operations 不需要特定的 dtype 来实现稳定性，而是接受多个输入
并要求 inputs 的 dtypes 匹配。如果所有输入都是 ，则运算在 中运行。如果任何输入为 ，
autocast 将所有输入强制转换为 并在 中运行运算。float16float16float32float32float32
addcdiv, , , , , , , , ,addcmulatan2bilinearcrossdotgrid_sampleindex_putscatter_addtensordot
此处未列出的一些操作（例如，像这样的二进制操作）原生会提升
inputs 的 intent 请求。如果输入是 和 的混合，则这些运算会运行并生成输出
无论是否启用了 Autocast。addfloat16float32float32float32
优先
（和 ） 的向后传递 （包装它）
可以生成在 中无法表示的渐变。在启用了自动广播的区域中，正向输入
可能是 ，这意味着向后渐变必须在 （自动转换 forward inputs to 没有帮助，因为该转换必须在 backward 中反转）。
因此，并在启用自动转换的区域中引发错误。float16float16float16float16float32binary_cross_entropyBCELoss
许多模型在二进制交叉熵层之前使用 sigmoid 层。
在这种情况下，请使用 或 合并两个图层。 并且可以安全地自动施法。binary_cross_entropy_with_logitsBCEWithLogits
XPU Op 特定行为（实验性）¶
以下列表描述了启用自动转换的区域中符合条件的操作的行为。
这些操作始终会进行自动转换，无论它们是作为 ，
作为函数或方法。如果函数在多个命名空间中公开，则
无论命名空间如何，它们都会进行自动转换。
下面未列出的 Ops 不会进行自动转换。它们在 由他们的输入定义。但是，自动转换可能仍会更改类型 如果未列出的操作位于自动转换操作的下游，则运行其中的未列出操作。
如果一个运算未列出，我们假设它在 中的数值稳定。
如果您认为未列出的操作在 中数值不稳定 ，
请提交 issue。float16float16
可以自动施放到
addbmm, , , , , , , , , , , , , , , , , , , ,addmmaddmvaddrbaddbmmbmmchain_matmulmulti_dotconv1dconv2dconv3dconv_transpose1dconv_transpose2dconv_transpose3dGRUCelllinearLSTMCellmatmulmmmvRNNCell
可以自动施放到
__pow__, , , , , , , , , , , , , , , , , , , , , , , , , , , ,__rdiv____rpow____rtruediv__binary_cross_entropy_with_logitscosine_embedding_losscosine_similaritycumsumdistexpgroup_normhinge_embedding_losskl_divl1_losslayer_normloglog_softmaxmargin_ranking_lossnll_lossnormalizepoisson_nll_losspowreciprocalrsqrtsoft_margin_losssoftmaxsoftminsumtriplet_margin_loss
提升到最宽输入类型的 XPU Ops¶
这些 operations 不需要特定的 dtype 来实现稳定性，而是接受多个输入
并要求 inputs 的 dtypes 匹配。如果所有输入都是 ，则运算在 中运行。如果任何输入为 ，
autocast 将所有输入强制转换为 并在 中运行运算。float16float16float32float32float32
bilinear, , , , ,crossgrid_sampleindex_putscatter_addtensordot
此处未列出的一些操作（例如，像这样的二进制操作）原生会提升
inputs 的 intent 请求。如果输入是 和 的混合，则这些运算会运行并生成输出
无论是否启用了 Autocast。addfloat16float32float32float32
特定于 CPU 操作的行为¶
以下列表描述了启用自动转换的区域中符合条件的操作的行为。
这些操作始终会进行自动转换，无论它们是作为 ，
作为函数或方法。如果函数在多个命名空间中公开，则
无论命名空间如何，它们都会进行自动转换。
下面未列出的 Ops 不会进行自动转换。它们在 由他们的输入定义。但是，自动转换可能仍会更改类型 如果未列出的操作位于自动转换操作的下游，则运行其中的未列出操作。
如果一个运算未列出，我们假设它在 中的数值稳定。
如果您认为未列出的操作在 中数值不稳定 ，
请提交 issue。 共享 的列表。bfloat16bfloat16float16bfloat16
可自动转换为
conv1d, , , , , , , , , , , , , , , , , , ,conv2dconv3dbmmmmlinalg_vecdotbaddbmmaddmmaddbmmlinearmatmul_convolutionconv_tbcmkldnn_rnn_layerconv_transpose1dconv_transpose2dconv_transpose3dpreluscaled_dot_product_attention_native_multi_head_attention
可自动转换为
avg_pool3d, , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , ,binary_cross_entropygrid_samplergrid_sampler_2d_grid_sampler_2d_cpu_fallbackgrid_sampler_3dpolarprodquantilenanquantilestftcdisttraceview_as_complexcholeskycholesky_inversecholesky_solveinverselu_solveorgqrinverseormqrpinversemax_pool3dmax_unpool2dmax_unpool3dadaptive_avg_pool3dreflection_pad1dreflection_pad2dreplication_pad1dreplication_pad2dreplication_pad3dmse_losscosine_embedding_lossnll_lossnll_loss2dhinge_embedding_losspoisson_nll_losscross_entropy_lossl1_losshuber_lossmargin_ranking_losssoft_margin_losstriplet_margin_lossmulti_margin_lossctc_losskl_divmultilabel_margin_lossbinary_cross_entropy_with_logitsfft_fftfft_ifftfft_fft2fft_ifft2fft_fftnfft_ifftnfft_rfftfft_irfftfft_rfft2fft_irfft2fft_rfftnfft_irfftnfft_hfftfft_ihfftlinalg_condlinalg_matrix_ranklinalg_solvelinalg_choleskylinalg_svdvalslinalg_eigvalslinalg_eigvalshlinalg_invlinalg_householder_productlinalg_tensorinvlinalg_tensorsolvefake_quantize_per_tensor_affinegeqrf_lu_with_infoqrsvdtriangular_solvefractional_max_pool2dfractional_max_pool3dadaptive_max_pool3dmultilabel_margin_loss_forwardlinalg_qrlinalg_cholesky_exlinalg_svdlinalg_eiglinalg_eighlinalg_lstsqlinalg_inv_ex
提升到最宽输入类型的 CPU Ops¶
这些 operations 不需要特定的 dtype 来实现稳定性，而是接受多个输入
并要求 inputs 的 dtypes 匹配。如果所有输入都是 ，则运算在 中运行。如果任何输入为 ，
autocast 将所有输入强制转换为 并在 中运行运算。bfloat16bfloat16float32float32float32
cat, ,stackindex_copy
此处未列出的一些操作（例如，像这样的二进制操作）原生会提升
inputs 的 intent 请求。如果输入是 和 的混合，则这些运算会运行并生成输出
无论是否启用了 Autocast。addbfloat16float32float32float32