混合精度自动包 - torch.cuda.amp¶

torch.cuda.amp 和 torch 为混合精度提供了便捷方法，其中一些操作使用 torch.float32 (float) 数据类型，而其他操作使用 torch.float16 (half)。一些操作，如线性层和卷积, 在 float16 中要快得多。其他操作，如归约操作，通常需要 float32 的动态范围。混合精度尝试将每个操作匹配到其合适的数据类型。

通常，“自动混合精度训练”使用 torch.autocast 和 torch.cuda.amp.GradScaler 一起，如在自动混合精度示例和自动混合精度示例中所示。但是，torch.autocast 和 GradScaler 是模块化的，并且可以根据需要单独使用。

自动类型转换 ¶

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

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

参见自动混合精度示例以了解在更复杂场景（例如梯度惩罚、多个模型/损失、自定义自动求导函数）中的使用方法（以及梯度缩放）。

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

class AutocastModel(nn.Module):
    ...
    @autocast()
    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 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 some tensors in default dtype (here assumed to be float32)
a_float32 = torch.rand((8, 8), device="cpu")
b_float32 = torch.rand((8, 8), device="cpu")
c_float32 = torch.rand((8, 8), device="cpu")
d_float32 = torch.rand((8, 8), device="cpu")

with autocast(dtype=torch.bfloat16, device_type="cpu"):
    # torch.mm is on autocast's list of ops that should run in bfloat16.
    # Inputs are float32, but the op runs in bfloat16 and produces bfloat16 output.
    # No manual casts are required.
    e_bfloat16 = torch.mm(a_float32, b_float32)
    # Also handles mixed input types
    f_bfloat16 = torch.mm(d_float32, e_bfloat16)

# After exiting autocast, calls f_float16.float() to use with d_float32
g_float32 = torch.mm(d_float32, f_bfloat16.float())

类型不匹配错误是在自动混合精度区域中的一个 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 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状态是线程本地的。如果你想在新线程中启用它，必须在该线程中调用上下文管理器或装饰器。这会影响torch.nn.DataParallel和torch.nn.parallel.DistributedDataParallel，当与每个进程多个GPU一起使用时（参见使用多个GPU）。

Parameters

device_type (字符串, 必填) – 使用 ‘cuda’ 或 ‘cpu’ 设备
已启用 (bool, 可选, 默认=True) – 是否在该区域中启用自动类型转换。
数据类型 (torch_dpython:type, 可选) – 是否使用 torch.float16 或 torch.bfloat16。
cache_enabled (bool, optional, default=True) – 是否启用 autocast 内部的权重缓存。

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, **kwargs)[source]¶

Helper decorator for forward methods of custom autograd functions (subclasses of torch.autograd.Function). See the example page for more detail.

Parameters: cast_inputs (torch.dtype 或 None，可选，默认=None) – 如果不是 None, 当 forward 在启用 autocast 的区域中运行时，会将传入的浮点 CUDA 张量转换为目标数据类型（非浮点张量不受影响），然后在禁用 autocast 的情况下执行 forward。如果 None，则 forward 的内部操作将使用当前的 autocast 状态执行。

注意

如果装饰的 forward 在未启用autocast的区域中被调用， custom_fwd 将不起作用，并且 cast_inputs 没有任何效果。

torch.cuda.amp.custom_bwd(bwd)[source]¶: 用于自定义autograd函数（torch.autograd.Function的子类）的反向方法的辅助装饰器。确保backward在与forward相同的autocast状态下执行。有关详细信息，请参阅示例页面。

梯度缩放 ¶

如果某个操作的前向传递有 float16 个输入，那么该操作的后向传递将产生 float16 个梯度。幅度较小的梯度值可能无法在 float16 中表示。这些值将被冲为零（“下溢”），因此相应参数的更新将会丢失。

为防止下溢，“梯度缩放”会将网络的损失乘以一个缩放因子，并对缩放后的损失执行反向传播。通过网络向后流动的梯度也会被相同的因子缩放。换句话说，梯度值的幅度更大，因此不会归零。

每个参数的梯度（.grad 属性）应在优化器更新参数之前进行解缩放，以便缩放因子不会干扰学习率。

class torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)[source]¶

get_backoff_factor()[source]¶: 返回一个包含缩放回退因子的 Python 浮点数。

get_growth_factor()[source]¶: 返回一个包含缩放增长因子的 Python 浮点数。

get_growth_interval()[source]¶: 返回一个包含增长间隔的 Python 整数。

get_scale()[source]¶: 返回一个包含当前缩放比例的 Python 浮点数，如果禁用缩放则返回 1.0。

警告

get_scale() 会引发 CPU-GPU 同步。

is_enabled()[source]¶: 返回一个布尔值，表示此实例是否已启用。

load_state_dict(state_dict)[source]¶

加载缩放器状态。如果此实例被禁用，则 load_state_dict() 是一个空操作。

Parameters: state_dict (dict) – 标量状态。应该是一个从对 state_dict() 的调用返回的对象。

scale(outputs)[source]¶

将张量或张量列表乘以（“缩放”）一个缩放因子。

返回缩放后的输出。如果此实例的 GradScaler 未启用，则输出将保持不变。

Parameters: 输出 (张量或 张量的可迭代对象) – 需要缩放的输出。

set_backoff_factor(new_factor)[source]¶

Parameters: new_scale (float) – 用作新缩放回退因子的值。

set_growth_factor(new_factor)[source]¶

Parameters: new_scale (float) – 用作新尺度增长因子的值。

set_growth_interval(new_interval)[source]¶

Parameters: new_interval (int) – 用作新增长间隔的值。

state_dict()[source]¶

返回缩放器的状态为 dict。它包含五个条目：

"scale" - 一个包含当前比例的 Python 浮点数
"growth_factor" - 一个包含当前增长因子的Python浮点数
"backoff_factor" - 一个包含当前回退因子的Python浮点数
"growth_interval" - 一个包含当前增长间隔的Python整数
"_growth_tracker" - 一个Python整数，表示最近连续未跳过的步骤数量。

如果此实例未启用，则返回一个空字典。

注意

如果你希望在特定迭代后保存缩放器的状态，state_dict() 应该在 update() 之后调用。

step(optimizer, *args, **kwargs)[source]¶

step() 执行以下两个操作：

内部调用 unscale_(optimizer)（除非在迭代的早期已经为 optimizer 明确调用了 unscale_()）。作为 unscale_() 的一部分，会检查梯度是否存在 inf/NaN。
如果没有发现无穷大/NaN梯度，则使用未缩放的梯度调用optimizer.step()。否则，将跳过optimizer.step()以避免损坏参数。

*args 和 **kwargs 被转发到 optimizer.step()。

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

Parameters

优化器 (torch.optim.Optimizer) – 应用梯度的优化器。
args – 任何参数。
kwargs – 任意关键字参数。

警告

目前不支持闭包的使用。

unscale_(optimizer)[source]¶

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

unscale_() 是可选的，适用于您需要在反向传递之间修改或检查梯度和 step() 之间的情况。如果 unscale_() 没有被显式调用，梯度将在 step() 期间自动缩放。

简单示例，使用 unscale_() 来启用未缩放梯度的裁剪：

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

Parameters: 优化器 (torch.optim.Optimizer) – 拥有需要取消缩放梯度的优化器。

注意

unscale_() 不会导致 CPU-GPU 同步。

警告

unscale_() 应该在每个优化器的每次 step() 调用中仅调用一次，并且仅在为该优化器分配的所有参数都累积了梯度之后。对于给定的优化器，在每次 step() 之间两次调用 unscale_() 将触发 RuntimeError。

警告

unscale_() 可能会在原地取消稀疏梯度的缩放，替换 .grad 属性。

update(new_scale=None)[source]¶

更新缩放因子。

如果任何优化器步骤被跳过，比例将乘以 backoff_factor 以减小它。如果连续发生 growth_interval 次未跳过的迭代，比例将乘以 growth_factor 以增加它。

Passing new_scale sets the new scale value manually. (new_scale is not used directly, it’s used to fill GradScaler’s internal scale tensor. So if new_scale was a tensor, later in-place changes to that tensor will not further affect the scale GradScaler uses internally.)

Parameters: new_scale (float 或 torch.cuda.FloatTensor, 可选, default=None) – 新的缩放因子。

警告

update() 应该只在迭代的最后调用，在为本次迭代使用的所有优化器调用了 scaler.step(optimizer) 之后。

自动混合精度操作参考文档 ¶

资格要求 ¶

只有 CUDA 操作符有资格进行自动精度转换。

运行在float64或非浮点数数据类型的运算不可用，并且无论是否启用了自动类型转换，都会在此数据类型上运行。

只有非原地操作和张量方法符合条件。在启用了自动类型转换的区域中，允许使用原地变体和显式提供out=...张量的调用，但它们不会经过自动类型转换。例如，在启用了自动类型转换的区域中，a.addmm(b, c)可以进行自动类型转换，但a.addmm_(b, c)和a.addmm(b, c, out=d)不能。为了获得最佳性能和稳定性，请在启用了自动类型转换的区域中优先使用非原地操作。

显式调用带有 dtype=... 参数的操作不符合要求，并将产生尊重 dtype 参数的输出。

特定操作的行为 ¶

以下列表描述了在自动类型转换启用区域中符合条件的操作的行为。这些操作始终会经历自动类型转换，无论它们是作为torch.nn.Module的一部分被调用，作为一个函数，还是作为一个torch.Tensor方法。如果函数在多个命名空间中暴露，它们都会经历自动类型转换，而与命名空间无关。

未在下面列出的操作不会经过自动类型转换。它们按照其输入定义的类型运行。然而，如果未列出的操作位于自动类型转换操作的下游，自动类型转换仍可能会更改这些操作的运行类型。

如果一个操作未列出，我们假设它在float16中数值稳定。如果你认为一个未列出的操作在float16中数值不稳定，请提交一个问题。

Ops that can autocast to `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

Ops that can autocast to `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

促进到最宽输入类型的运算 ¶

这些操作不需要特定的数据类型来保证稳定性，但需要多个输入且输入的数据类型必须匹配。如果所有的输入都是 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 可以安全地自动转换。