高级 Amp 用法¶

GANs¶

GANs 是下面几个主题的有趣综合。综合示例正在构建中。

梯度裁剪¶

Amp 将优化器的 param_groups 直接拥有的参数称为“主参数”。

这些主参数可能与 model.parameters() 完全或部分不同。例如，使用 opt_level="O2"，amp.initialize 将大多数模型参数转换为 FP16，为每个新 FP16 模型参数在模型外部创建一个 FP32 主参数，并更新优化器的 param_groups 以指向这些 FP32 参数。

优化器的 param_groups 拥有的主参数也可能与模型参数完全一致，这通常适用于 opt_levels O0、O1 和 O3。

在所有情况下，正确的做法是裁剪保证由 优化器的 param_groups 拥有的参数的梯度，而不是通过 model.parameters() 检索到的参数的梯度。

此外，如果 Amp 使用损失缩放，则必须在取消缩放梯度后（在退出 amp.scale_loss 上下文管理器期间发生）裁剪梯度。

以下模式对于任何 opt_level 应该是正确的

with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
    # Gradients are unscaled during context manager exit.
# Now it's safe to clip.  Replace
# torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
# with
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_norm)
# or
torch.nn.utils.clip_grad_value_(amp.master_params(optimizer), max_)

请注意使用实用函数 amp.master_params(optimizer)，它返回一个生成器表达式，该表达式迭代优化器的 param_groups 中的参数。

另请注意，调用 clip_grad_norm_(amp.master_params(optimizer), max_norm) 而不是，而不是除了，clip_grad_norm_(model.parameters(), max_norm)。

自定义/用户定义的 autograd 函数¶

注册用户函数的旧 Amp API 仍然被认为是正确的。必须在调用 amp.initialize 之前注册函数。

强制特定层/函数使用所需的类型¶

我仍在努力寻找一种通用的暴露方式，它不需要用户端代码在不同的 opt-levels 之间产生差异。

多个模型/优化器/损失¶

使用多个模型/优化器初始化¶

amp.initialize 的优化器参数可以是单个优化器或优化器列表，只要您接受的输出具有相同的类型。同样，model 参数可以是单个模型或模型列表，只要接受的输出匹配。以下调用都是合法的

model, optim = amp.initialize(model, optim,...)
model, [optim0, optim1] = amp.initialize(model, [optim0, optim1],...)
[model0, model1], optim = amp.initialize([model0, model1], optim,...)
[model0, model1], [optim0, optim1] = amp.initialize([model0, model1], [optim0, optim1],...)

使用多个优化器进行反向传播¶

每当您调用反向传播时，amp.scale_loss 上下文管理器必须接收 所有拥有任何参数的优化器，当前的反向传播正在为这些参数创建梯度。 即使每个优化器只拥有部分，而不是全部将要接收梯度的参数，也是如此。

如果对于给定的反向传播，只有一个优化器将接收梯度的参数，您可以直接将该优化器传递给 amp.scale_loss。否则，您必须传递将接收梯度的参数的优化器列表。包含 3 个损失和 2 个优化器的示例

# loss0 accumulates gradients only into params owned by optim0:
with amp.scale_loss(loss0, optim0) as scaled_loss:
    scaled_loss.backward()

# loss1 accumulates gradients only into params owned by optim1:
with amp.scale_loss(loss1, optim1) as scaled_loss:
    scaled_loss.backward()

# loss2 accumulates gradients into some params owned by optim0
# and some params owned by optim1
with amp.scale_loss(loss2, [optim0, optim1]) as scaled_loss:
    scaled_loss.backward()

可以选择让 Amp 为每个损失使用不同的损失缩放器¶

默认情况下，Amp 维护一个全局损失缩放器，该缩放器将用于所有反向传播（所有调用 with amp.scale_loss(...)）。使用全局损失缩放器不需要额外的 amp.initialize 或 amp.scale_loss 参数。上面带有多个优化器/反向传播的代码片段在幕后使用单个全局损失缩放器，它们应该“正常工作”。

但是，您可以选择告诉 Amp 为每个损失维护一个损失缩放器，这使 Amp 具有更高的数值灵活性。这是通过将 num_losses 参数提供给 amp.initialize 来实现的（这告诉 Amp 您计划调用多少个反向传播，因此 Amp 应该创建多少个损失缩放器），然后将 loss_id 参数提供给每个反向传播（这告诉 Amp 用于此特定反向传播的损失缩放器）

model, [optim0, optim1] = amp.initialize(model, [optim0, optim1], ..., num_losses=3)

with amp.scale_loss(loss0, optim0, loss_id=0) as scaled_loss:
    scaled_loss.backward()

with amp.scale_loss(loss1, optim1, loss_id=1) as scaled_loss:
    scaled_loss.backward()

with amp.scale_loss(loss2, [optim0, optim1], loss_id=2) as scaled_loss:
    scaled_loss.backward()

num_losses 和 loss_id 应纯粹根据损失/反向传播的集合来指定。使用多个优化器，或将单个或多个优化器与每个反向传播相关联是无关的。

跨迭代的梯度累积¶

以下应该“正常工作”，并正确容纳多个模型/优化器/损失，以及通过上述说明进行梯度裁剪

# If your intent is to simulate a larger batch size using gradient accumulation,
# you can divide the loss by the number of accumulation iterations (so that gradients
# will be averaged over that many iterations):
loss = loss/iters_to_accumulate

with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

# Every iters_to_accumulate iterations, call step() and reset gradients:
if iter%iters_to_accumulate == 0:
    # Gradient clipping if desired:
    # torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), max_norm)
    optimizer.step()
    optimizer.zero_grad()

作为一种小的性能优化，您可以将 delay_unscale=True 传递给 amp.scale_loss，直到您准备好 step()。只有当您确定知道自己在做什么时，才应尝试 delay_unscale=True，因为与梯度裁剪和多个模型/优化器/损失的交互可能会变得棘手。

if iter%iters_to_accumulate == 0:
    # Every iters_to_accumulate iterations, unscale and step
    with amp.scale_loss(loss, optimizer) as scaled_loss:
        scaled_loss.backward()
    optimizer.step()
    optimizer.zero_grad()
else:
    # Otherwise, accumulate gradients, don't unscale or step.
    with amp.scale_loss(loss, optimizer, delay_unscale=True) as scaled_loss:
        scaled_loss.backward()

自定义数据批次类型¶

Amp 的目的是无论 opt_level 如何，您都无需手动转换输入数据。 Amp 通过修补任何模型的 forward 方法来适当地为 opt_level 转换传入数据来实现此目的。但是，要转换传入数据，Amp 需要知道如何。修补的 forward 将识别并转换浮点张量（像 IntTensors 这样的非浮点张量不会被触及）和浮点张量的 Python 容器。但是，如果您将张量包装在一个自定义类中，则转换逻辑不知道如何穿透坚硬的自定义外壳以访问和转换其中的多汁张量。您需要通过为其分配一个接受 torch.dtype （例如，torch.float16 或 torch.float32）并返回转换为 dtype 的自定义批次实例的 to 方法来告诉 Amp 如何转换自定义批次类。修补的 forward 检查您的 to 方法是否存在，并将使用 opt_level 的正确类型调用它。

示例

class CustomData(object):
    def __init__(self):
        self.tensor = torch.cuda.FloatTensor([1,2,3])

    def to(self, dtype):
        self.tensor = self.tensor.to(dtype)
        return self

警告

Amp 也会转发 numpy ndarrays 而不转换它们。如果您将输入数据作为原始的、未包装的 ndarray 发送，然后在 model.forward 中使用它来创建一个张量，则此张量的类型将不取决于 opt_level，并且可能是正确的，也可能是不正确的。鼓励用户尽可能传递可转换的数据输入（张量、张量集合或具有 to 方法的自定义类）。

注意

Amp 不会为您在任何张量上调用 .cuda()。 Amp 假设您的原始脚本已经设置为根据需要将张量从主机移动到设备。