• 文档 >
  • 自动混合精度示例
快捷方式

自动混合精度示例

通常,“自动混合精度训练”是指使用torch.autocasttorch.amp.GradScaler 共同进行训练。

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

torch.amp.GradScaler 的实例有助于方便地执行梯度缩放步骤。梯度缩放通过最小化梯度下溢来改善具有 float16(默认在 CUDA 和 XPU 上)梯度的网络的收敛性,如此处所述。

torch.autocasttorch.amp.GradScaler 是模块化的。在下面的示例中,每个都按其单独的文档建议使用。

(此处的示例仅供说明。有关可运行的演练,请参阅自动混合精度食谱。)

典型的混合精度训练

# Creates model and optimizer in default precision
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# Creates a GradScaler once at the beginning of training.
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # Runs the forward pass with autocasting.
        with autocast(device_type='cuda', dtype=torch.float16):
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
        # Backward passes under autocast are not recommended.
        # Backward ops run in the same dtype autocast chose for corresponding forward ops.
        scaler.scale(loss).backward()

        # scaler.step() first unscales the gradients of the optimizer's assigned params.
        # If these gradients do not contain infs or NaNs, optimizer.step() is then called,
        # otherwise, optimizer.step() is skipped.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        scaler.update()

使用未缩放的梯度

scaler.scale(loss).backward() 生成的所有梯度都已缩放。如果您希望在 backward()scaler.step(optimizer) 之间修改或检查参数的 .grad 属性,则应首先对其进行反缩放。例如,梯度裁剪会操作一组梯度,使其全局范数(参见torch.nn.utils.clip_grad_norm_())或最大幅度(参见torch.nn.utils.clip_grad_value_()<=<=某个用户设定的阈值。如果您尝试在进行反缩放的情况下进行裁剪,则梯度的范数/最大幅度也会被缩放,因此您请求的阈值(原本是针对未缩放梯度的阈值)将无效。

scaler.unscale_(optimizer) 会反缩放 optimizer 所分配的参数持有的梯度。如果您的模型或模型包含分配给另一个优化器(例如 optimizer2)的其他参数,您可以分别调用 scaler.unscale_(optimizer2) 来反缩放这些参数的梯度。

梯度裁剪

在裁剪之前调用 scaler.unscale_(optimizer) 使您能够像往常一样裁剪未缩放的梯度

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()

        # Unscales the gradients of optimizer's assigned params in-place
        scaler.unscale_(optimizer)

        # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        # optimizer's gradients are already unscaled, so scaler.step does not unscale them,
        # although it still skips optimizer.step() if the gradients contain infs or NaNs.
        scaler.step(optimizer)

        # Updates the scale for next iteration.
        scaler.update()

scaler 记录了此迭代中是否已为该优化器调用了 scaler.unscale_(optimizer),因此 scaler.step(optimizer) 知道在(内部)调用 optimizer.step() 之前不要冗余地反缩放梯度。

警告

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

使用缩放的梯度

梯度累积

梯度累积会在一个有效批大小为 batch_per_iter * iters_to_accumulate(如果分布式则为 * num_procs)上累加梯度。缩放因子应该针对有效批大小进行校准,这意味着 inf/NaN 检查、如果发现 inf/NaN 梯度则跳过步骤,以及缩放因子更新应该以有效批大小为粒度进行。此外,在累积给定有效批大小的梯度期间,梯度应该保持缩放状态,并且缩放因子应该保持不变。如果在累积完成之前梯度被取消缩放(或缩放因子发生变化),则下一个反向传播将把缩放的梯度添加到未缩放的梯度(或缩放因子不同的梯度),之后无法恢复累积的未缩放梯度 step 必须应用。

因此,如果您想 unscale_ 梯度(例如,允许裁剪未缩放的梯度),请在 step 之前调用 unscale_,在即将进行的 step 的所有(缩放的)梯度累积之后。此外,仅在您为完整有效批大小调用了 step 的迭代结束时调用 update

scaler = GradScaler()

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

        # Accumulates scaled gradients.
        scaler.scale(loss).backward()

        if (i + 1) % iters_to_accumulate == 0:
            # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)

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

梯度惩罚

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

这是一个没有梯度缩放或自动转换的普通 L2 惩罚示例。

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)

        # Creates gradients
        grad_params = torch.autograd.grad(outputs=loss,
                                          inputs=model.parameters(),
                                          create_graph=True)

        # Computes the penalty term and adds it to the loss
        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()

        # clip gradients here, if desired

        optimizer.step()

要实现带梯度缩放的梯度惩罚,传递给 torch.autograd.grad()outputs 张量应该被缩放。因此,生成的梯度将被缩放,并且应该在组合以创建惩罚值之前取消缩放。

此外,惩罚项计算是前向传播的一部分,因此应该在 autocast 上下文中。

以下是相同 L2 惩罚的外观。

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)

        # Scales the loss for autograd.grad's backward pass, producing scaled_grad_params
        scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss),
                                                 inputs=model.parameters(),
                                                 create_graph=True)

        # Creates unscaled grad_params before computing the penalty. scaled_grad_params are
        # not owned by any optimizer, so ordinary division is used instead of scaler.unscale_:
        inv_scale = 1./scaler.get_scale()
        grad_params = [p * inv_scale for p in scaled_grad_params]

        # Computes the penalty term and adds it to the loss
        with autocast(device_type='cuda', dtype=torch.float16):
            grad_norm = 0
            for grad in grad_params:
                grad_norm += grad.pow(2).sum()
            grad_norm = grad_norm.sqrt()
            loss = loss + grad_norm

        # Applies scaling to the backward call as usual.
        # Accumulates leaf gradients that are correctly scaled.
        scaler.scale(loss).backward()

        # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)

        # step() and update() proceed as usual.
        scaler.step(optimizer)
        scaler.update()

使用多个模型、损失和优化器

如果您的网络有多个损失,则必须分别对每个损失调用 scaler.scale。如果您的网络有多个优化器,您可以分别对其中任何一个调用 scaler.unscale_,并且必须分别对每个优化器调用 scaler.step

但是,scaler.update 应该只调用一次,在本次迭代中使用的所有优化器都执行了 step 之后。

scaler = torch.amp.GradScaler()

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

        # (retain_graph here is unrelated to amp, it's present because in this
        # example, both backward() calls share some sections of graph.)
        scaler.scale(loss0).backward(retain_graph=True)
        scaler.scale(loss1).backward()

        # You can choose which optimizers receive explicit unscaling, if you
        # want to inspect or modify the gradients of the params they own.
        scaler.unscale_(optimizer0)

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

        scaler.update()

每个优化器都会检查其梯度中是否存在 inf/NaN,并独立决定是否跳过该步骤。这可能导致一个优化器跳过该步骤,而另一个优化器不跳过。由于跳过步骤很少发生(每几百次迭代),因此不应阻碍收敛。如果您在将梯度缩放添加到多优化器模型后观察到收敛性变差,请报告错误。

使用多个 GPU

此处描述的问题仅影响 autocastGradScaler 的用法保持不变。

单个进程中的 DataParallel

即使 torch.nn.DataParallel 生成线程在每个设备上运行前向传播。autocast 状态在每个线程中都得到传播,并且以下内容将起作用。

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

# Sets autocast in the main thread
with autocast(device_type='cuda', dtype=torch.float16):
    # dp_model's internal threads will autocast.
    output = dp_model(input)
    # loss_fn also autocast
    loss = loss_fn(output)

DistributedDataParallel,每个进程一个 GPU

torch.nn.parallel.DistributedDataParallel 的文档建议每个进程一个 GPU 以获得最佳性能。在这种情况下,DistributedDataParallel 在内部不会生成线程,因此 autocastGradScaler 的用法不受影响。

DistributedDataParallel,每个进程多个 GPU

在这里,torch.nn.parallel.DistributedDataParallel 可能会生成一个辅助线程在每个设备上运行前向传播,就像 torch.nn.DataParallel 一样。修复方法相同:将 autocast 应用为模型的 forward 方法的一部分,以确保它在辅助线程中启用。

Autocast 和自定义 Autograd 函数

如果您的网络使用 自定义 Autograd 函数torch.autograd.Function 的子类),如果任何函数,则需要进行更改以实现与 autocast 的兼容性。

  • 接受多个浮点张量输入,

  • 包装任何可自动转换的操作(请参阅 Autocast 操作参考),或

  • 需要特定的 dtype(例如,如果它包装了仅为 dtype 编译的 CUDA 扩展)。

在所有情况下,如果您正在导入该函数并且无法更改其定义,则一个安全的回退方法是在发生错误的任何使用点禁用 autocast 并强制以 float32(或 dtype)执行。

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

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

具有多个输入或可自动转换操作的函数

custom_fwdcustom_bwd(无参数)分别应用于 forwardbackward。这些确保 forward 使用当前的 autocast 状态执行,并且 backward 使用与 forward 相同的 autocast 状态执行(这可以防止类型不匹配错误)。

class MyMM(torch.autograd.Function):
    @staticmethod
    @custom_fwd
    def forward(ctx, a, b):
        ctx.save_for_backward(a, b)
        return a.mm(b)
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        a, b = ctx.saved_tensors
        return grad.mm(b.t()), a.t().mm(grad)

现在,可以在任何地方调用 MyMM,而无需禁用 autocast 或手动转换输入。

mymm = MyMM.apply

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

需要特定 dtype 的函数

考虑一个需要 torch.float32 输入的自定义函数。将 custom_fwd(device_type='cuda', cast_inputs=torch.float32) 应用于 forward,并将 custom_bwd(device_type='cuda') 应用于 backward。如果 forward 在启用 autocast 的区域中运行,则装饰器会将浮点张量输入转换为 float32,并在参数 device_type 指定的设备上(本例中为 CUDA),并在 forwardbackward 期间局部禁用 autocast。

class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @custom_fwd(device_type='cuda', cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        ...
        return fwd_output
    @staticmethod
    @custom_bwd(device_type='cuda')
    def backward(ctx, grad):
        ...

现在,可以在任何地方调用 MyFloat32Func,而无需手动禁用 autocast 或转换输入。

func = MyFloat32Func.apply

with autocast(device_type='cuda', dtype=torch.float16):
    # func will run in float32, regardless of the surrounding autocast state
    output = func(input)

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并获得问题的解答

查看资源