apex.amp¶

本页面记录了 Amp（自动混合精度）的更新 API，该工具只需 3 行 Python 代码即可实现 Tensor Core 加速训练。

在 Github 页面上可以找到一个演示良好实践的可运行、全面的 Imagenet 示例。

GAN 是一个棘手的案例，许多人都有此需求。一个全面的 DCGAN 示例正在建设中。

如果您已经根据下面的说明实现了 Amp，但其行为不如预期，请查阅Amp 高级用法，看看是否有与您的用例匹配的主题。如果仍然没有帮助，请提交一个问题。

`opt_level`s 和属性¶

Amp 允许用户轻松尝试不同的纯精度和混合精度模式。通过选择“优化级别”或 opt_level 来选择常用的默认模式；每个 opt_level 都会建立一组属性，这些属性决定了 Amp 如何实现纯精度或混合精度训练。通过将特定属性的值直接传递给 amp.initialize，可以实现对给定 opt_level 行为的更细粒度控制。这些手动指定的值会覆盖由 opt_level 建立的默认值。

示例

# Declare model and optimizer as usual, with default (FP32) precision
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

# Allow Amp to perform casts as required by the opt_level
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
...
# loss.backward() becomes:
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
...

用户不应该手动将其模型或数据转换为 .half()，无论选择了哪个 opt_level 或属性。Amp 的目标是让用户从现有的默认（FP32）脚本开始，添加与 Amp API 对应的三行代码，然后开始混合精度训练。Amp 也可以被禁用，在这种情况下，原始脚本将完全按照之前的方式运行。这样一来，遵循 Amp API 没有风险，并且有很大的潜在性能提升。

注意

因为除了调用 amp.initialize 之外，永远不需要手动转换模型或输入数据，因此遵循新 API 的脚本可以在不同的 opt-level 之间切换，而无需进行任何其他更改。

属性¶

目前，控制纯精度或混合精度训练的底层属性如下

cast_model_type: 将模型的参数和缓冲区转换为所需的类型。
patch_torch_functions: 修补所有 Torch 函数和 Tensor 方法，以便在 FP16 中执行 Tensor Core 友好的操作（例如 GEMM 和卷积），并在 FP32 中执行任何受益于 FP32 精度的操作。
keep_batchnorm_fp32: 为了提高精度并启用 cudnn batchnorm（这可以提高性能），通常有益于即使模型的其余部分是 FP16，也保持 batchnorm 权重在 FP32 中。
master_weights: 维护 FP32 主权重以配合任何 FP16 模型权重。FP32 主权重由优化器进行步进，以提高精度并捕获微小梯度。
loss_scale: 如果 loss_scale 是浮点值，则将此值用作静态（固定）损失缩放。如果 loss_scale 是字符串 "dynamic"，则随时间自适应调整损失缩放。动态损失缩放调整由 Amp 自动执行。

同样，您通常不需要手动指定这些属性。而是选择一个 opt_level，它会为您设置它们。选择 opt_level 后，您可以选择性地将属性 kwargs 作为手动覆盖传递。

如果您尝试覆盖对于选定的 opt_level 没有意义的属性，Amp 将会引发错误并提供解释。例如，选择 opt_level="O1" 并结合覆盖 master_weights=True 是没有意义的。O1 是在 Torch 函数周围插入类型转换，而不是模型权重。数据、激活和权重在流经修补后的函数时会即时地进行就地转换。因此，模型权重本身可以（并且应该）保持 FP32，并且没有必要维护单独的 FP32 主权重。

`opt_level`s¶

可识别的 opt_levels 有 "O0", "O1", "O2" 和 "O3"。

O0 和 O3 不是真正的混合精度，但它们分别用于建立精度和速度基线，非常有用。

O1 和 O2 是混合精度的不同实现。可以都尝试一下，看看哪种能为您的模型带来最佳的速度提升和精度。

`O0`: FP32 训练¶

您的传入模型应该是 FP32 格式，所以这可能是一个空操作。O0 对于建立精度基线很有用。

O0 设置的默认属性
cast_model_type=torch.float32
patch_torch_functions=False
keep_batchnorm_fp32=None （实际上，“不适用”，一切都是 FP32）
master_weights=False
loss_scale=1.0

`O1`: 混合精度 (推荐用于典型场景)¶

修补所有 Torch 函数和 Tensor 方法，根据白名单-黑名单模型转换它们的输入。白名单操作（例如，Tensor Core 友好的操作，如 GEMM 和卷积）在 FP16 中执行。受益于 FP32 精度的黑名单操作（例如 softmax）在 FP32 中执行。O1 还会使用动态损失缩放，除非被覆盖。

O1 设置的默认属性
cast_model_type=None （不适用）
patch_torch_functions=True
keep_batchnorm_fp32=None （同样不适用，所有模型权重保持 FP32）
master_weights=None （不适用，模型权重保持 FP32）
loss_scale="dynamic"

`O2`: “近似 FP16” 混合精度¶

O2 将模型权重转换为 FP16，修补模型的 forward 方法以将输入数据转换为 FP16，保持 batchnorm 为 FP32，维护 FP32 主权重，更新优化器的 param_groups，以便 optimizer.step() 直接作用于 FP32 权重（如有必要，之后进行 FP32 主权重到 FP16 模型权重的拷贝），并实现动态损失缩放（除非被覆盖）。与 O1 不同，O2 不会修补 Torch 函数或 Tensor 方法。

O2 设置的默认属性
cast_model_type=torch.float16
patch_torch_functions=False
keep_batchnorm_fp32=True
master_weights=True
loss_scale="dynamic"

`O3`: FP16 训练¶

O3 可能无法达到真正的混合精度选项 O1 和 O2 的稳定性。但是，它对于为您的模型建立速度基线很有用，可以用来比较 O1 和 O2 的性能。如果您的模型使用了批量归一化，为了达到“光速”，您可以尝试 O3 并额外覆盖属性 keep_batchnorm_fp32=True（这会启用 cudnn batchnorm，如前所述）。

O3 设置的默认属性
cast_model_type=torch.float16
patch_torch_functions=False
keep_batchnorm_fp32=False
master_weights=False
loss_scale=1.0

统一 API¶

apex.amp.initialize(models, optimizers=None, enabled=True, opt_level='O1', cast_model_type=None, patch_torch_functions=None, keep_batchnorm_fp32=None, master_weights=None, loss_scale=None, cast_model_outputs=None, num_losses=1, verbosity=1, min_loss_scale=None, max_loss_scale=16777216.0)[source]¶

根据选择的 opt_level 和任何被覆盖的属性，初始化您的模型、优化器以及 Torch tensor 和 functional 命名空间。

amp.initialize 应该在您完成模型和优化器的构建之后调用，但在您将模型送入任何 DistributedDataParallel 包装器之前调用。请参阅 Imagenet 示例中的分布式训练。

目前，amp.initialize 应该只调用一次，尽管它可以处理任意数量的模型和优化器（请参阅相应的Amp 高级用法主题）。如果您认为您的用例需要多次调用 amp.initialize，请告知我们。

任何不是 None 的属性关键字参数将被解释为手动覆盖。

为了避免不得不重写脚本中的其他任何内容，请将返回的模型/优化器命名为替换传入的模型/优化器，如下面的代码示例所示。

参数

models (torch.nn.Module 或 list of torch.nn.Modules) – 要修改/转换的模型。
optimizers (可选, torch.optim.Optimizer 或 list of torch.optim.Optimizers) – 要修改/转换的优化器。训练时必需，推理时可选。
enabled (bool, 可选, default=True) – 如果为 False，所有 Amp 调用都将成为空操作，您的脚本将像没有 Amp 一样运行。
opt_level (str, 可选, default="O1") – 纯精度或混合精度优化级别。接受的值为 “O0”, “O1”, “O2” 和 “O3”，上面已详细解释。
cast_model_type (torch.dtype, 可选, default=None) – 可选的属性覆盖，见上文。
patch_torch_functions (bool, 可选, default=None) – 可选的属性覆盖。
keep_batchnorm_fp32 (bool 或 str, 可选, default=None) – 可选的属性覆盖。如果作为字符串传递，必须是字符串 “True” 或 “False”。
master_weights (bool, 可选, default=None) – 可选的属性覆盖。
loss_scale (float 或 str, 可选, default=None) – 可选的属性覆盖。如果作为字符串传递，必须是表示数字的字符串，例如 “128.0”，或者是字符串 “dynamic”。
cast_model_outputs (torch.dpython:type, 可选, default=None) – 一个选项，用于确保您的模型的输出始终转换为特定类型，无论 opt_level 如何。
num_losses (int, 可选, default=1) – 一个选项，用于提前告知 Amp 您计划使用多少个损失/反向传播。与 amp.scale_loss 的 loss_id 参数结合使用时，使 Amp 能够为每个损失/反向传播使用不同的损失缩放，这可以提高稳定性。有关多个优化器的示例，请参阅Amp 高级用法中的“多个模型/优化器/损失函数”。如果 num_losses 保留为 1，Amp 仍然支持多个损失/反向传播，但会为所有这些损失使用一个全局损失缩放。
verbosity (int, default=1) – 设置为 0 以禁止 Amp 相关输出。
min_loss_scale (float, default=None) – 设置动态损失缩放可以选择的损失缩放值的下限。默认值 None 表示没有施加下限。如果未使用动态损失缩放，则忽略 min_loss_scale。
max_loss_scale (float, default=2.**24) – 设置动态损失缩放可以选择的损失缩放值的上限。如果未使用动态损失缩放，则忽略 max_loss_scale。

返回值

根据 opt_level 修改后的模型和优化器。如果 models 或 optimizers 参数是列表，则相应的返回值也将是列表。

允许的调用方式

model, optim = amp.initialize(model, optim,...)
model, [optim1, optim2] = amp.initialize(model, [optim1, optim2],...)
[model1, model2], optim = amp.initialize([model1, model2], optim,...)
[model1, model2], [optim1, optim2] = amp.initialize([model1, model2], [optim1, optim2],...)

# This is not an exhaustive list of the cross product of options that are possible,
# just a set of examples.
model, optim = amp.initialize(model, optim, opt_level="O0")
model, optim = amp.initialize(model, optim, opt_level="O0", loss_scale="dynamic"|128.0|"128.0")

model, optim = amp.initialize(model, optim, opt_level="O1") # uses "loss_scale="dynamic" default
model, optim = amp.initialize(model, optim, opt_level="O1", loss_scale=128.0|"128.0")

model, optim = amp.initialize(model, optim, opt_level="O2") # uses "loss_scale="dynamic" default
model, optim = amp.initialize(model, optim, opt_level="O2", loss_scale=128.0|"128.0")
model, optim = amp.initialize(model, optim, opt_level="O2", keep_batchnorm_fp32=True|False|"True"|"False")

model, optim = amp.initialize(model, optim, opt_level="O3") # uses loss_scale=1.0 default
model, optim = amp.initialize(model, optim, opt_level="O3", loss_scale="dynamic"|128.0|"128.0")
model, optim = amp.initialize(model, optim, opt_level="O3", keep_batchnorm_fp32=True|False|"True"|"False")

Imagenet 示例演示了各种 opt_levels 和覆盖的实际使用。

apex.amp.scale_loss(loss, optimizers, loss_id=0, model=None, delay_unscale=False, delay_overflow_check=False)[source]¶

在上下文管理器进入时，创建 scaled_loss = (loss.float())*current loss scale。scaled_loss 被 yielding 出来，以便用户可以调用 scaled_loss.backward()

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

在上下文管理器退出时（如果 delay_unscale=False），梯度会被检查是否有 infs/NaNs，然后进行反缩放，以便可以调用 optimizer.step()。

注意

如果 Amp 使用显式的 FP32 主参数（这是 opt_level=O2 的默认设置，也可以通过向 amp.initialize 提供 master_weights=True 手动启用），任何 FP16 梯度在反缩放之前都会被复制到 FP32 主梯度。optimizer.step() 然后会将反缩放后的主梯度应用于主参数。

警告

如果 Amp 使用显式的 FP32 主参数，则只有 FP32 主梯度会被反缩放。任何 FP16 模型参数的直接 .grad 属性在上下文管理器退出后仍保持缩放。这种细微之处会影响梯度裁剪。有关最佳实践，请参阅Amp 高级用法中的“梯度裁剪”。

参数

loss (Tensor) – 通常是标量 Tensor。上下文管理器 yield 的 scaled_loss 只是 loss.float()*loss_scale，因此原则上 loss 可以有多个元素，只要您在上下文管理器主体内适当地对 scaled_loss 调用 backward()。
optimizers – 当前反向传播正在为其创建梯度的所有优化器。必须是之前调用 amp.initialize 返回的优化器或优化器列表。有关与多个优化器一起使用的示例，请参阅Amp 高级用法中的“多个模型/优化器/损失函数”。
loss_id (int, 可选, default=0) – 与 amp.initialize 的 num_losses 参数结合使用时，使 Amp 能够为每个损失使用不同的损失缩放。loss_id 必须是一个介于 0 和 num_losses 之间的整数，它告诉 Amp 当前反向传播使用的是哪个损失。有关示例，请参阅Amp 高级用法中的“多个模型/优化器/损失函数”。如果未指定 loss_id，Amp 将对本次反向传播使用默认的全局损失缩放器。
model (torch.nn.Module, 可选, default=None) – 当前未使用，保留用于未来优化。
delay_unscale (bool, 可选, default=False) – delay_unscale 永远不是必需的，并且强烈建议使用默认值 False。如果为 True，Amp 在上下文管理器退出时将不会反缩放梯度或执行模型到主梯度的拷贝。delay_unscale=True 是一个微不足道的性能优化技巧，可能会导致奇怪的问题（尤其是在有多个模型/优化器/损失函数时），所以仅在您了解其作用时使用。在Amp 高级用法中的“迭代间的梯度累积”下举例说明了可以使用（但非必须）这种情况。

警告

如果在给定反向传播中 delay_unscale 为 True，则在上下文管理器退出后还不能调用 optimizer.step()，必须等待另一个后续的反向传播上下文管理器调用，其中 delay_unscale 保持为 False。

apex.amp.master_params(optimizer)[source]¶

生成器表达式，遍历 optimizer 所拥有的参数。

参数: optimizer – 之前从 amp.initialize 返回的优化器。

检查点¶

为了正确保存和加载您的 amp 训练，我们引入了 amp.state_dict()，它包含所有 loss_scaler 及其相应的未跳过步数，以及 amp.load_state_dict() 来恢复这些属性。

为了获得逐位精度，我们建议采用以下工作流程

# Initialization
opt_level = 'O1'
model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)

# Train your model
...

# Save checkpoint
checkpoint = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    'amp': amp.state_dict()
}
torch.save(checkpoint, 'amp_checkpoint.pt')
...

# Restore
model = ...
optimizer = ...
checkpoint = torch.load('amp_checkpoint.pt')

model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
amp.load_state_dict(checkpoint['amp'])

# Continue training
...

请注意，我们建议使用相同的 opt_level 来恢复模型。另请注意，我们建议在调用 amp.initialize 之后调用 load_state_dict 方法。

高级用例¶

统一的 Amp API 支持迭代间的梯度累积、每次迭代的多次反向传播、多个模型/优化器、自定义/用户定义的 autograd 函数以及自定义数据批处理类。梯度裁剪和 GAN 也需要特殊处理，但这种处理对于不同的 opt_level 不需要改变。更多详情可在此处找到

Amp 高级用法

旧 API 用户迁移指南¶

我们强烈建议迁移到新的 Amp API，因为它更通用、更易于使用且面向未来。原始的 FP16_Optimizer 和旧的“Amp” API 已被弃用，并可能随时被移除。

旧“Amp” API 用户¶

在新 API 中，opt-level O1 执行与旧版“Amp”相同的 Torch 命名空间修补。但是，新 API 允许静态或动态损失缩放，而旧 API 只允许动态损失缩放。

在新 API 中，旧的调用 amp_handle = amp.init() 和返回的 amp_handle 不再暴露或需要。新的 amp.initialize() 承担了 amp.init() 的职责（以及更多）。因此，任何现有的 amp_handle = amp.init() 调用都应该被删除。

以前通过 amp_handle 暴露的函数现在是通过 amp 模块访问的自由函数。

反向传播上下文管理器必须相应地改变

# old API
with amp_handle.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
->
# new API
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

目前，已弃用的“Amp” API 文档仍可在 Github README 中找到：https://github.com/NVIDIA/apex/tree/master/apex/amp。旧 API 中用于标注用户函数以特定精度运行的调用在新 API 中仍然有效。

旧 FP16_Optimizer 用户¶

opt-level O2 等同于带有 dynamic_loss_scale=True 的 FP16_Optimizer。再次强调，反向传播必须改为统一版本

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

FP16_Optimizer 的一个令人烦恼之处在于，用户必须手动将模型转换为 half 精度（通过在其上调用 .half()，或使用 apex.fp16_utils 中的函数或模块包装器），并且还要手动对输入数据调用 .half()。在新 API 中，这些都不需要。无论您选择哪个 --opt-level，您都可以并且应该直接构建您的模型并以默认的 FP32 格式传递输入数据。新的 Amp API 将在 model, optimizer = amp.initialize(model, optimizer, opt_level=....) 期间，根据 --opt-level 和任何被覆盖的标志进行正确的转换。浮点输入数据可以是 FP32 或 FP16，但您最好让它保持 FP16，因为 amp.initialize 返回的 model 将会修补其 forward 方法以适当转换输入数据。