自动混合精度软件包 - torch.amp

为混合精度提供便捷的方法， 其中，某些操作使用 （） 数据类型，而其他操作 使用低精度浮点数据类型 （）：（） 或 .一些运算，如线性层和卷积、 在 中要快得多。其他操作（如缩减）通常需要动态 的范围。混合精度尝试将每个运算与其适当的数据类型匹配。torch.float32floatlower_precision_fptorch.float16halftorch.bfloat16lower_precision_fpfloat32

通常，“自动混合精度训练”的数据类型为 uses 和 together，如 CUDA Automatic Mixed Precision 示例和 CUDA Automatic Mixed Precision 配方所示。 但是，并且是模块化的，如果需要，可以单独使用。 如 CPU 示例部分所示，在 CPU 上使用 “自动混合精度训练/推理” 数据类型仅使用 .torch.float16torch.bfloat16

对于 CUDA 和 CPU，API 也单独提供：

• torch.autocast("cuda", args...)等效于 。torch.cuda.amp.autocast(args...)

• torch.autocast("cpu", args...)等效于 。对于 CPU，目前仅支持较低精度的浮点数据类型 of。torch.cpu.amp.autocast(args...)torch.bfloat16

自动投射

类 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 autocast():
        output = model(input)
        loss = loss_fn(output, target)

    # Exits the context manager before backward()
    loss.backward()
    optimizer.step()

有关用法（以及梯度缩放），请参阅 CUDA 自动混合精度示例 在更复杂的场景中（例如，梯度惩罚、多个模型/损失、自定义 Autograd 函数）。

也可以用作装饰器，例如，在你的模型的方法上：forward

class AutocastModel(nn.Module):
    ...
    @autocast()
    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 autocast():
    # 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 autocast():
    e_float16 = torch.mm(a_float32, b_float32)
    with autocast(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 设备

• enabled （bool， optional） – 是否应在区域中启用自动转换。 违约：True

• dtype （torch_dtype，可选） – 是使用 torch.float16 还是 torch.bfloat16。

• cache_enabled （bool， optional） – 是否应启用自动转换中的权重缓存。 违约：True

类 torch.cuda.amp。autocast（enabled=True， dtype=torch.float16， cache_enabled=True）

请参阅 。 等效于torch.cuda.amp.autocast(args...)torch.autocast("cuda", args...)

torch.cuda.amp。custom_fwd（fwd=无， *， cast_inputs=无）

自定义 autograd 函数（的子类）方法的辅助装饰器 。有关更多详细信息，请参阅示例页面。forward

参数

cast_inputs （ 或 None， optional， default=None） – 如果不是 ， 在启用了自动广播的区域运行时，将转换传入的 浮点 CUDA 张量添加到目标 dtype 中（非浮点张量不受影响）， 然后在禁用自动转换的情况下执行。 如果 ， 的内部操作以当前自动转换状态执行。NoneforwardforwardNoneforward

注意

如果在启用了自动转换的区域之外调用 decored，则为 no-op，并且没有效果。forwardcast_inputs

torch.cuda.amp。custom_bwd（BWD）

自定义 autograd 函数（的子类）的向后方法的辅助装饰器 。 确保 以与 相同的自动转换状态执行。 有关更多详细信息，请参阅示例页面。backwardforward

类 torch.cpu.amp 中。autocast（enabled=True， dtype=torch.bfloat16， cache_enabled=True）

请参阅 。 等效于torch.cpu.amp.autocast(args...)torch.autocast("cpu", args...)

梯度缩放

如果特定运算的正向传递有输入，则 该运算将产生梯度。 具有较小量级的梯度值可能无法在 中表示。 这些值将刷新为零（“下溢”），因此相应参数的更新将丢失。float16float16float16

为了防止下溢，“梯度缩放”将网络的损失乘以比例因子，并且 对缩放的损失调用向后传递。通过网络向后流动的梯度是 然后按相同的因子进行缩放。换句话说，梯度值具有更大的幅度， 因此，它们不会刷新为零。

每个参数的 gradient （ attribute） 应在优化器之前取消缩放 更新参数，以便 Scale Factor 不会干扰学习率。.grad

类 torch.cuda.amp。GradScaler（init_scale=65536.0， growth_factor=2.0， backoff_factor=0.5， growth_interval=2000， enabled=True）

get_backoff_factor（）

返回包含刻度回退因子的 Python 浮点数。

get_growth_factor（）

返回包含 scale 增长因子的 Python 浮点数。

get_growth_interval（）

返回包含增长区间的 Python int。

get_scale（）

返回包含当前缩放的 Python 浮点数，如果禁用缩放，则返回 1.0。

警告

导致 CPU-GPU 同步。

is_enabled（

返回一个 bool，指示是否启用此实例。

load_state_dict（state_dict）

加载缩放器状态。如果此实例已禁用，则为 no-op。

参数

state_dict （dict） – 缩放器状态。应为从调用 .

scale（outputs）

将张量或张量列表乘以比例因子。

返回缩放后的输出。如果未启用此实例，则返回输出 未修改。

参数

outputs （Tensor 或 Tensors 的可迭代对象） – 要缩放的输出。

set_backoff_factor（new_factor）

参数

new_scale （float） （float） （ （float） （浮点数） ） – 用作新比例回退因子的值。

set_growth_factor（new_factor）

参数

new_scale （float） （ （float） （浮点数） ） – 用作新比例增长因子的值。

set_growth_interval（new_interval）

参数

new_interval （int） – 用作新增长区间的值。

state_dict（）

将缩放器的状态返回为 .它包含 5 个条目：

• "scale"- 包含当前刻度的 Python 浮点数

• "growth_factor"- 包含当前增长因子的 Python 浮点数

• "backoff_factor"- 包含当前回退因子的 Python 浮点数

• "growth_interval"- 包含当前增长区间的 Python int

• "_growth_tracker"- 一个 Python int，其中包含最近连续未跳过的步骤数。

如果未启用此实例，则返回空 dict。

注意

如果您希望在特定迭代后对缩放器的状态进行检查点操作，则应在之后调用。

step（optimizer， *args， **kwargs）

执行以下两个操作：

1. 内部调用 （除非在迭代的早期明确调用）。作为 的一部分，将检查 infs/NaN 的梯度。unscale_(optimizer)optimizer

2. 如果未找到 inf/NaN 梯度，则使用未缩放的 梯度。否则，将跳过以避免损坏参数。optimizer.step()optimizer.step()

*args并转发到 。**kwargsoptimizer.step()

返回 的返回值 。optimizer.step(*args, **kwargs)

参数

• optimizer （torch.optim.Optimizer） – 应用梯度的优化器。

• args – 任何参数。

• kwargs – 任何关键字参数。

警告

目前不支持使用 Closure。

unscale_（优化器）

将优化器的梯度张量除以比例因子（“取消缩放”）。

是可选的，适用于需要修改或检查向后传递 和 之间的梯度的情况。 如果未显式调用，则渐变将在 期间自动取消缩放。

简单示例，用于启用未缩放渐变的裁剪：

...
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(optimizer)
scaler.update()

参数

optimizer （torch.optim.Optimizer） – 拥有要未缩放的梯度的优化器。

注意

不会导致 CPU-GPU 同步。

警告

每个优化器每次调用只应调用一次， 并且仅在该优化器的 assign parameters 的所有梯度都已累积之后。 在每次调用之间为给定优化器调用两次会触发 RuntimeError。

警告

可能会取消稀疏渐变的缩放比例，从而替换该属性。.grad

update（new_scale=None）

更新比例因子。

如果跳过了任何优化器步骤，则 scale 乘以以将其减少。如果未跳过的迭代连续发生，则 将 scale 乘以增加它。backoff_factorgrowth_intervalgrowth_factor

传递 （Passpass） 将手动设置新的缩放值。（ 不是 直接使用，用于填充 GradScaler 的内部 scale 张量。因此，如果是一个张量，则以后对该张量的就地更改将不会进一步 影响 GradScaler 内部使用的比例。new_scalenew_scalenew_scale

参数

new_scale （float or ， optional， default=None） – 新比例因子。torch.cuda.FloatTensor

警告

应该只在迭代结束时调用，在 has 被调用。scaler.step(optimizer)

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

特定于 CPU 操作的行为

以下列表描述了启用自动转换的区域中符合条件的操作的行为。 这些操作始终会进行自动转换，无论它们是作为 ， 作为函数或方法。如果函数在多个命名空间中公开，则 无论命名空间如何，它们都会进行自动转换。

下面未列出的 Ops 不会进行自动转换。它们在 由他们的输入定义。但是，自动转换可能仍会更改类型 如果未列出的操作位于自动转换操作的下游，则运行其中的未列出操作。

如果一个运算未列出，我们假设它在 中的数值稳定。 如果您认为未列出的操作在 中数值不稳定 ， 请提交 issue。bfloat16bfloat16

可自动转换为

conv1d, , , , , , , , , ,conv2dconv3dbmmmmbaddbmmaddmmaddbmmlinearmatmul_convolution

可自动转换为

conv_transpose1d, , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , , ,conv_transpose2dconv_transpose3davg_pool3dbinary_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_lossctc_losskl_divmultilabel_margin_lossfft_fftfft_ifftfft_fft2fft_ifft2fft_fftnfft_ifftnfft_rfftfft_irfftfft_rfft2fft_irfft2fft_rfftnfft_irfftnfft_hfftfft_ihfftlinalg_matrix_normlinalg_condlinalg_matrix_ranklinalg_solvelinalg_choleskylinalg_svdvalslinalg_eigvalslinalg_eigvalshlinalg_invlinalg_householder_productlinalg_tensorinvlinalg_tensorsolvefake_quantize_per_tensor_affineeiggeqrflstsq_lu_with_infoqrsolvesvdsymeigtriangular_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