•
pytorch-doc-2.5中文文档
首页
/
pytorch文档
/
自动混合精度示例

自动混合精度示例

通常，“自动混合精度训练”意味着使用 torch.autocast 和 torch.amp.GradScaler 一起进行训练。

torch.autocast 的实例可以为选定的区域启用自动混合精度。自动混合精度会自动选择操作的精度，以提高性能同时保持准确性。

torch.amp.GradScaler 的实例有助于方便地执行梯度缩放步骤。梯度缩放通过最小化梯度下溢（即梯度值过小导致数值不稳定）来改善具有 float16（默认在 CUDA 和 XPU 上）梯度的网络的收敛性，具体解释见 这里。

torch.autocast 和 torch.amp.GradScaler 是模块化设计的。在下面的示例中，每个组件都按照其各自的文档建议使用。

（以下示例仅供参考。请参阅 自动混合精度教程 获取可运行的指南。）

典型的混合精度训练

## [处理未缩放的梯度](#id3)

所有由 `scaler.scale(loss).backward()` 产生的梯度都是缩放过的。如果您希望在 `backward()` 和 `scaler.step(optimizer)` 之间修改或检查参数的 `.grad` 属性，应首先将梯度恢复到未缩放的状态。例如，梯度裁剪会调整一组梯度，使其全局范数（参见 [`torch.nn.utils.clip_grad_norm_()`](../generated/torch.nn.utils.clip_grad_norm_.html#torch.nn.utils.clip_grad_norm_)）或最大值（参见 [`torch.nn.utils.clip_grad_value_()`](../generated/torch.nn.utils.clip_grad_value_.html#torch.nn.utils.clip_grad_value_)）小于等于某个用户设定的阈值。如果您尝试在不取消缩放的情况下进行裁剪，梯度的范数或最大值也会被缩放，因此您设定的阈值（原本是针对未缩放的梯度的）将不再有效。

`scaler.unscale_(optimizer)` 会将 `optimizer` 分配的参数所持有的梯度恢复到未缩放的状态。如果您的模型中还有其他参数被分配给了另一个优化器（例如 `optimizer2`），您可以单独调用 `scaler.unscale_(optimizer2)` 来将这些参数的梯度恢复到未缩放的状态。

### [梯度裁剪](#id4)

在裁剪之前调用 `scaler.unscale_(optimizer)` 可以让你像往常一样裁剪未缩放的梯度：

```python
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()

        # 对优化器分配的参数的梯度进行原地去缩放（即直接修改原梯度值）
        scaler.unscale_(optimizer)

        # 因为优化器分配的参数的梯度已经被去缩放，所以可以像往常一样进行裁剪：
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        # 优化器的梯度已经取消缩放，因此 `scaler.step` 不会再次取消缩放它们，
        # 尽管如此，如果梯度包含 inf 或 NaN，`scaler.step` 仍然会跳过 `optimizer.step()`。
        scaler.step(optimizer)

        # 更新用于下一次迭代的缩放因子。
        scaler.update()

scaler 记录了在当前迭代中已经为该优化器调用了 scaler.unscale_(optimizer)，因此 scaler.step(optimizer) 在内部调用 optimizer.step() 之前不会重复取消缩放梯度。

```

警告

每个优化器在每个 `step` 调用中只能调用一次 `unscale_`，并且只能在该优化器分配的参数的所有梯度都累积完毕后调用。对于给定的优化器，在每次 `step` 之间调用两次 `unscale_` 会触发 `RuntimeError`。

## [处理缩放梯度](#id5)

### [梯度累积](#id6)

梯度累积会在一个有效的批次大小为 `batch_per_iter * iters_to_accumulate`（如果是分布式训练，则再乘以 `num_procs`）上累加梯度。缩放因子应针对有效的批次进行校准，这包括进行 inf/NaN 检查、如果发现 inf/NaN 梯度则跳过步骤，并且应在有效批次粒度上更新缩放因子。此外，在累积给定有效批次的梯度时，梯度应保持缩放状态，缩放因子应保持不变。如果在累积完成之前取消缩放梯度（或缩放因子发生变化），下一个反向传播将把缩放的梯度加到未缩放的梯度（或由不同因子缩放的梯度）上，之后将无法恢复累积的未缩放梯度，`step` 必须应用。

因此，如果您希望对梯度进行 `unscale_`（例如，允许裁剪未缩放的梯度），请在 `step` 之前立即调用 `unscale_`，在即将进行的 `step` 中所有（已缩放的）梯度都已累积之后。此外，仅在调用了 `step` 的迭代结束时调用 `update`，以确保处理了一个完整的有效批次：

plain scaler = GradScaler()

for epoch in epochs: for i, (input, target) in enumerate(data): with autocast(devicetype='cuda', dtype=torch.float16): output = model(input) loss = lossfn(output, target) loss = loss / iterstoaccumulate

    # 累积缩放的梯度。
    scaler.scale(loss).backward()

    if (i + 1) % iters_to_accumulate == 0:
        # 如果需要，可以在这里调用 `unscale_`（例如，允许裁剪未缩放的梯度）

        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

### [梯度惩罚](#id7)

梯度惩罚的实现通常使用 `torch.autograd.grad()` 创建梯度，将这些梯度组合起来创建惩罚值，并将惩罚值添加到损失中。

这是一个普通的 L2 惩罚示例，不涉及梯度缩放或自动混合精度：

plain for epoch in epochs: for input, target in data: optimizer.zerograd() output = model(input) loss = lossfn(output, target)

    # 计算梯度
    grad_params = torch.autograd.grad(outputs=loss,
                                      inputs=model.parameters(),
                                      create_graph=True)

    # 计算惩罚项并将其加到损失中
    grad_norm = 0
    for grad in grad_params:
        grad_norm += grad.pow(2).sum()
    grad_norm = grad_norm.sqrt()
    loss = loss + grad_norm

    loss.backward()

    # 如果需要，可以在这里对梯度进行裁剪

    optimizer.step()

为了实现带有梯度缩放的梯度惩罚，传递给 `torch.autograd.grad()` 的 `outputs` 张量应该进行缩放。因此，生成的梯度也会被缩放，在组合成惩罚值之前需要取消缩放。

此外，惩罚项的计算是前向传播的一部分，因此应该在 `autocast` 环境中进行。

以下是如何处理相同 L2 惩罚项的示例。

## [处理多个模型、损失和优化器](#id8)

如果您的网络包含多个损失函数，您需要对每个损失函数分别调用 `scaler.scale`。如果您的网络包含多个优化器，您可以对其中任何一个单独调用 `scaler.unscale_`，并且必须对每个优化器分别调用 `scaler.step`。

但是，`scaler.update` 只应在所有优化器在本次迭代中使用过后调用一次：

plain scaler = torch.amp.GradScaler()

for epoch in epochs: for input, target in data: optimizer0.zerograd() optimizer1.zerograd() with autocast(devicetype='cuda', dtype=torch.float16): output0 = model0(input) output1 = model1(input) loss0 = lossfn(2 * output0 + 3 * output1, target) loss1 = loss_fn(3 * output0 - 5 * output1, target)

    # （这里使用 `retain_graph` 与自动混合精度无关，它存在是因为在这个例子中，
    # 两个 `backward()` 调用共享了一些图的部分。）
    scaler.scale(loss0).backward(retain_graph=True)
    scaler.scale(loss1).backward()

    # 您可以选择对哪些优化器进行显式的反缩放，以便检查或修改这些优化器拥有的参数的梯度。
    scaler.unscale_(optimizer0)

    scaler.step(optimizer0)
    scaler.step(optimizer1)

    scaler.update()

每个优化器都会检查其梯度是否存在无穷大（infs）或非数字（NaNs），并独立决定是否跳过该步。这可能导致一个优化器跳过该步，而另一个优化器则不跳过。由于跳过步骤的情况很少发生（每几百次迭代一次），这不应影响收敛。如果您在向多优化器模型添加梯度缩放后发现收敛效果不佳，请报告一个 bug。

## [使用多个 GPU](#id9)

这些问题仅影响 `autocast`。`GradScaler` 的使用方式保持不变。

### [单进程中的 DataParallel](#id10)

即使 [`torch.nn.DataParallel`](../generated/torch.nn.DataParallel.html#torch.nn.DataParallel) 会在每个设备上启动线程进行前向传播，autocast 状态也会在每个线程中传播，以下代码将正常工作：

plain model = MyModel() dp_model = nn.DataParallel(model)

在主线程中启用 autocast

with autocast(devicetype='cuda', dtype=torch.float16): # dpmodel 的内部线程将启用 autocast output = dpmodel(input) # lossfn 也启用了 autocast loss = loss_fn(output)

### [分布式数据并行，每个进程一个GPU](#id11)

[`torch.nn.parallel.DistributedDataParallel`](../generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel) 的文档建议每个进程使用一个GPU以获得最佳性能。在这种情况下，`DistributedDataParallel` 不会启动内部线程，因此 `自动类型转换` 和 `GradScaler` 的使用不受影响。

### [分布式数据并行，每个进程多个GPU](#id12)

在这种情况下，`DistributedDataParallel` 可能在每个设备上启动一个辅助线程来运行前向传递，类似于 `torch.nn.DataParallel`。[解决方案相同](#amp-dataparallel)：在模型的 `forward` 方法中应用 `自动类型转换`，以确保在辅助线程中启用它。

## [自动类型转换和自定义自动梯度函数](#id13)

如果您的网络使用了 [自定义自动梯度函数](extending.html#extending-autograd)（即 `torch.autograd.Function` 的子类），则需要进行一些更改以确保与自动类型转换兼容，特别是当任何函数

* 接受多个浮点 Tensor 输入，

* 包装任何支持自动类型转换的操作（参见 [自动类型转换操作参考](../amp.html#autocast-op-reference)），或

* 需要特定的数据类型 `dtype`（例如，如果它包装了仅编译为 `dtype` 的 [CUDA 扩展](https://pytorch.org/tutorials/advanced/cpp_extension.html)）。

在所有情况下，如果您正在导入该函数并且无法更改其定义，一个安全的回退方案是在出现错误时禁用自动转换并强制使用 `float32`（或特定的 `dtype`）：

plain with autocast(devicetype='cuda', dtype=torch.float16): … with autocast(devicetype='cuda', dtype=torch.float16, enabled=False): output = imported_function(input1.float(), input2.float())

如果您是该函数的作者（或可以修改其定义），更好的解决方案是使用 `torch.amp.custom_fwd()` 和 `torch.amp.custom_bwd()` 装饰器，如下面的相关示例所示。

### [具有多个输入或支持自动类型转换操作的函数](#id14)

应用 [`custom_fwd`](../amp.html#torch.amp.custom_fwd) 和 [`custom_bwd`](../amp.html#torch.amp.custom_bwd)（不带参数）到 `forward` 和 `backward` 方法。这些装饰器确保 `forward` 在当前的自动类型转换状态下执行，并且 `backward` 也在相同的自动类型转换状态下执行（这可以防止类型不匹配错误）：

python class MyMM(torch.autograd.Function): @staticmethod @customfwd def forward(ctx, a, b): ctx.saveforbackward(a, b) return a.mm(b) @staticmethod @custombwd def backward(ctx, grad): a, b = ctx.saved_tensors return grad.mm(b.t()), a.t().mm(grad)

现在可以在任何地方调用 `MyMM`，而无需禁用自动类型转换或手动转换输入数据：

python mymm = MyMM.apply

with autocast(device_type='cuda', dtype=torch.float16): output = mymm(input1, input2)

### [需要特定 `dtype` 的函数](#id15)

考虑一个需要 `torch.float32` 类型输入的自定义函数。将 [`custom_fwd(device_type='cuda', cast_inputs=torch.float32)`](../amp.html#torch.amp.custom_fwd) 应用于 `forward`，并将 [`custom_bwd(device_type='cuda')`](../amp.html#torch.amp.custom_bwd) 应用于 `backward`。如果 `forward` 在启用了自动类型转换的区域运行，装饰器会将浮点数 Tensor 输入转换为 `float32`，并在指定的设备（例如 CUDA）上局部禁用自动类型转换，同时在 `forward` 和 `backward` 中也保持这种状态。

plain class MyFloat32Func(torch.autograd.Function): @staticmethod @customfwd(devicetype='cuda', castinputs=torch.float32) def forward(ctx, input): ctx.saveforbackward(input) … return fwdoutput @staticmethod @custombwd(devicetype='cuda') def backward(ctx, grad): …

现在可以在任何地方调用 `MyFloat32Func`，而无需手动禁用自动类型转换或手动转换输入：

plain func = MyFloat32Func.apply

with autocast(device_type='cuda', dtype=torch.float16): # 无论外部的自动类型转换设置如何，func 都将以 float32 运行 output = func(input)

```

本页目录
© 2024 modoc.cn
沪ICP备18008978号-4