自动混合精度软件包 - 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

自动投射

class （device_type， dtype=None， enabled=True， cache_enabled=None）torch.autocast

的实例充当上下文管理器或装饰器，这些 允许脚本的区域以混合精度运行。

在这些区域中，运算在 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(TestModel, self).__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 （string， required） – 是使用 'cuda' 还是 'cpu' 设备

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

• dtype （torch_dpython：type， optional） – 是使用 torch.float16 还是 torch.bfloat16。

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

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

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

torch.cuda.amp.custom_fwd(fwd=None， **kwargs）

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

参数

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

注意

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

torch.cuda.amp.custom_bwd(BWD）

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

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

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

梯度缩放

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

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

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

类 （init_scale=65536.0， growth_factor=2.0， backoff_factor=0.5， growth_interval=2000， 已启用=真）torch.cuda.amp.GradScaler

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 （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(优化器， *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=无）

更新比例因子。

如果跳过了任何优化器步骤，则 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_solvematrix_rankorgqrinverseormqrpinversemax_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