apex.fp16_utils

此子模块包含旨在简化 NVIDIA 提出的混合精度训练方法的实用工具，该方法发表在 在 Parallel Forall 上的文章 和 GTC 2018 会议的以下议题中：使用混合精度训练神经网络：理论与实践 以及 使用混合精度训练神经网络：实际示例。对于 Pytorch 用户，特别推荐实际示例部分。

演示 apex.fp16_utils 的完整可运行 Python 脚本可以在 Github 页面找到

简单的 FP16_Optimizer 演示

使用 Imagenet 进行分布式混合精度训练

使用 word_language_model 进行混合精度训练

主参数和损失缩放的自动管理

类 apex.fp16_utils.FP16_Optimizer(init_optimizer, static_loss_scale=1.0, dynamic_loss_scale=False, dynamic_loss_args=None, verbose=True)

FP16_Optimizer 旨在封装现有的 PyTorch 优化器，并以对用户透明的方式管理静态或动态损失缩放以及主权重。对于标准用法，只需修改两行代码：创建 FP16_Optimizer 实例，以及修改对 backward 的调用。

示例

model = torch.nn.Linear(D_in, D_out).cuda().half()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
# Name the FP16_Optimizer instance to replace the existing optimizer
# (recommended but not required):
optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
...
# loss.backward() becomes:
optimizer.backward(loss)
...

动态损失缩放示例

...
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
                           # optional arg to control dynamic loss scaling behavior
                           # dynamic_loss_args={'scale_window' : 500})
                           # Usually, dynamic_loss_args is not necessary.

参数

• init_optimizer (torch.optim.optimizer) – 使用需要优化的参数创建的现有优化器。在内部，FP16_Optimizer 会将传入的优化器的 fp16 参数（如果有）替换为从原始参数复制而来的 fp32 主参数。FP16_Optimizer 还会存储对原始 fp16 参数的引用，并在每次 step 调用结束时，从 fp32 主副本更新这些 fp16 参数。

• static_loss_scale (float, 可选, 默认值=1.0) – 内部用于缩放模型计算的梯度的损失缩放值。任何 fp16 梯度都会被复制到 fp32，然后在应用于 fp32 主参数之前进行下缩放，因此 static_loss_scale 不应影响学习率。

• dynamic_loss_scale (bool, 可选, 默认值=False) – 使用动态损失缩放。如果为 True，这将覆盖任何 static_loss_scale 选项。

• dynamic_loss_args (dict, 可选, 默认值=None) – 将转发到内部 LossScaler 实例构造函数的 kwargs 字典。此字典的键必须与 LossScaler 构造函数接受的 kwargs 匹配。如果未指定 dynamic_loss_args，则将使用 LossScaler 的默认值。

• verbose (bool, 可选, 默认值=True) – 默认情况下，FP16_Optimizer 的构造函数会打印其正在接收的参数和参数组，作为健全性检查。如果这变得令人烦恼（例如对于大型模型），可以通过传入 verbose=False 来禁用。在动态损失缩放期间调整损失缩放时，verbose=False 不会禁用打印。

init_optimizer 应以通常方式构造。推荐 (尽管不是强制要求) 将新构造的 FP16_Optimizer 实例命名为替换 init_optimizer 的名称，原因有二：首先，这意味着文件中后续对同一名称的引用无需更改。其次，FP16_Optimizer 保留（作为实现细节）修改 init_optimizer 的权利。如果您确实为新的 FP16_Optimizer 实例选择了唯一的名称，您应该只使用这个新实例，因为原有的优化器可能不再按预期工作。

init_optimizer 可以是任何 PyTorch 优化器。它可以包含 fp16 和 fp32 参数的混合，这些参数组织成任意数量的 param_groups，具有不同的超参数。FP16_Optimizer 的构造函数将接收这些 param_groups 并记住它们。

调用

loss.backward()

必须替换为

optimizer.backward(loss)

因为 FP16_Optimizer 需要拥有反向传播过程的控制权，以实现损失缩放以及向主梯度复制。

注意

损失缩放，无论是静态还是动态，与学习率是正交的，因为梯度在应用之前会下缩放。这意味着调整损失缩放，或使用动态损失缩放，不应需要重新调整学习率或任何其他超参数。

高级选项

闭包: FP16_Optimizer 可以封装一个接受闭包的 PyTorch 优化器。请参阅 step 的文档字符串。

梯度裁剪: 使用 clip_master_grads。

多个损失: 如果您的模型从多个损失累积梯度，可以通过向 backward 提供 update_master_grads=False 来提高效率。请参阅 backward 的文档字符串。

手动调整损失缩放: 当前的损失缩放值可以通过以下方式获取或设置

print(optimizer.loss_scale)
optimizer.loss_scale = new_loss_scale

对于静态损失缩放，随时间手动调整损失缩放是一个合理的做法。在后续 epoch 中，梯度可能变得更小，并且可能需要更高的损失缩放，这类似于学习率调度。动态损失缩放更微妙（参阅 DynamicLossScaler），在这种情况下，不建议手动调整损失缩放。

多 GPU 训练: 如果封装的 init_optimizer 是由封装在 Pytorch DistributedDataParallel 或 Apex DistributedDataParallel 中的模型创建的，FP16_Optimizer 应该仍然按预期工作。

backward(loss, update_master_grads=True, retain_graph=False)

backward 执行以下概念步骤

1. fp32_loss = loss.float() (参见下面的第一个注意)

2. scaled_loss = fp32_loss*loss_scale

3. scaled_loss.backward()，它将缩放后的梯度累积到模型叶节点的 .grad 属性中（根据您模型的定义方式，这些叶节点可能是 fp16、fp32 或混合类型）。

4. 然后将 fp16 梯度复制到主参数的 .grad 属性中（参见第二个注意），这些属性保证是 fp32 类型的。

5. 最后，主梯度除以损失缩放值 (loss_scale)。

通过这种方式，在调用 backward 后，主参数拥有新的梯度，并且可以调用 step 了。

注意

backward 在应用损失缩放之前，内部会将损失转换为 fp32。如果用户提供了 fp16 损失值，这提供了一些额外的防溢出安全。然而，为了最大限度地防止溢出，用户应在将其提供给 backward 之前，以 fp32 计算损失准则（如 MSE、交叉熵等）。

警告

在调用 backward 后，模型叶节点中的梯度通常不应被视为有效，因为它们可能已经被缩放（对于动态损失缩放，缩放因子可能会随时间变化）。如果用户希望在调用 backward 后检查梯度，只有主梯度应被视为有效。可以通过 inspect_master_grad_data() 获取这些梯度。

参数

• loss – 用户模型输出的损失值。loss 可以是 float 或 half 类型（但请参阅上面的第一个注意）。

• update_master_grads (bool, 可选, 默认值=True) – 在此调用中将 fp16 梯度复制到 fp32 梯度的选项。将其设置为 False，用户可以延迟复制，这在一次迭代中对多个损失调用 backward 时，有助于消除冗余的 fp16->fp32 梯度复制。如果设置为 False，用户将负责在调用 step 之前调用 update_master_grads。

• retain_graph (bool, 可选, 默认值=False) – 将通常的 retain_graph=True 选项转发到对 loss.backward 的内部调用。如果正在使用 retain_graph 在调用 optimizer.step 之前累积来自多次反向传播的梯度值，也建议传入 update_master_grads=False (参见下面的示例)。

示例

# Ordinary operation:
optimizer.backward(loss)

# Naive operation with multiple losses (technically valid, but less efficient):
# fp32 grads will be correct after the second call,  but
# the first call incurs an unnecessary fp16->fp32 grad copy.
optimizer.backward(loss1)
optimizer.backward(loss2)

# More efficient way to handle multiple losses:
# The fp16->fp32 grad copy is delayed until fp16 grads from all
# losses have been accumulated.
optimizer.backward(loss1, update_master_grads=False)
optimizer.backward(loss2, update_master_grads=False)
optimizer.update_master_grads()

clip_master_grads(max_norm, norm_type=2)

通过 torch.nn.utils.clip_grad_norm 裁剪 fp32 主梯度。

参数

• max_norm (float 或 int) – 梯度的最大范数

• norm_type (float 或 int) – 使用的 p-范数类型。对于无穷范数，可以是 'inf'。

返回

当前 fp32 梯度的总范数 (视为一个单一向量)。

警告

如果最近计算的 fp16 梯度溢出（即如果 self.overflow 为 True），则返回 -1。

inspect_master_grad_data()

在使用 FP16_Optimizer 运行时，模型 fp16 叶节点的 .grad 属性不应被视为真实值，因为它们可能已经被缩放。在调用 fp16_optimizer_obj.backward(loss) 后，如果没有遇到溢出，fp32 主参数的 .grad 属性将包含正确除以损失缩放值的有效梯度。然而，由于 FP16_Optimizer 会展平部分参数，访问它们可能不太直观。inspect_master_grad_data 允许以与其关联的模型叶节点相对应的形状查看这些梯度。

返回

列表的列表（每个参数组对应一个列表）。每个参数组的列表是属于该组的 fp32 主参数的 .grad.data 属性的列表。

load_state_dict(state_dict)

加载由之前调用 state_dict() 创建的 state_dict。如果 fp16_optimizer_instance 是由某个 init_optimizer 构造的，而其参数又来自 model，则期望用户在调用 fp16_optimizer_instance.load_state_dict() 之前调用 model.load_state_dict()。

示例

model = torch.nn.Linear(D_in, D_out).cuda().half()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
...
checkpoint = torch.load("saved.pth")
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])

state_dict()

返回一个包含此 FP16_Optimizer 实例当前状态的字典。此字典包含 FP16_Optimizer 的属性，以及所包含的 PyTorch 优化器的 state_dict。示例

checkpoint = {}
checkpoint['model'] = model.state_dict()
checkpoint['optimizer'] = optimizer.state_dict()
torch.save(checkpoint, "saved.pth")

step(closure=None)

如果没有提供闭包，应在调用 fp16_optimizer_obj.backward(loss) 后调用 step。step 使用提供给 FP16_Optimizer 构造函数的优化器更新参数的 fp32 主副本，然后将更新后的 fp32 参数复制到 FP16_Optimizer 构造函数最初引用的 fp16 参数中，这样用户可以立即使用其模型运行另一次前向传播。

如果提供了闭包，可以在没有先调用 backward(loss) 的情况下调用 step。这种控制流程与包含闭包的 普通 PyTorch 优化器用法 完全相同。然而，用户应注意确保闭包内任何对 loss.backward() 的调用已被替换为 fp16_optimizer_obj.backward(loss)。

参数

closure (可选) – 将提供给最初传递给 FP16_Optimizer 构造函数的底层优化器的闭包。闭包应在 FP16_Optimizer 对象上调用 zero_grad()，计算损失，调用 backward(loss)，并返回损失。

带闭包的示例

# optimizer is assumed to be an FP16_Optimizer object, previously constructed from an
# existing pytorch optimizer.
for input, target in dataset:
    def closure():
        optimizer.zero_grad()
        output = model(input)
        loss = loss_fn(output, target)
        # loss.backward() becomes:
        optimizer.backward(loss)
        return loss
    optimizer.step(closure)

警告

目前，使用闭包调用 step 与动态损失缩放不兼容。

update_master_grads()

从存储的对 fp16 参数的引用中复制 .grad 属性，到由优化器直接更新的 fp32 主参数的 .grad 属性。update_master_grads 仅当调用 fp16_optimizer_obj.backward 时设置了 update_master_grads=False 才需要调用。

zero_grad(set_grads_to_None=False)

将 fp32 和 fp16 参数的梯度清零。

类 apex.fp16_utils.LossScaler(scale=1)

管理静态损失缩放的类。此类旨在与 FP16_Optimizer 交互，不应由用户直接操作。

通过 FP16_Optimizer 构造函数中的 static_loss_scale 参数启用 LossScaler 的使用。

参数

scale (float, 可选, default=1.0) – 损失缩放比例。

class apex.fp16_utils.DynamicLossScaler(init_scale=4294967296, scale_factor=2.0, scale_window=1000)

管理动态损失缩放的类。建议间接使用DynamicLossScaler，方法是将dynamic_loss_scale=True提供给FP16_Optimizer的构造函数。然而，理解DynamicLossScaler的工作原理非常重要，因为可以使用FP16_Optimizer构造函数的dynamic_loss_args参数更改默认选项。

损失缩放旨在解决在使用fp16网络进行长时间训练时遇到的梯度下溢问题。动态损失缩放首先尝试一个非常高的损失缩放比例。讽刺的是，这可能会导致梯度溢出。如果遇到梯度溢出，DynamicLossScaler会通知FP16_Optimizer发生了溢出。FP16_Optimizer随后跳过这次迭代/minibatch的更新步骤，并且DynamicLossScaler将损失缩放比例调整为较低的值。如果连续多次迭代没有检测到梯度溢出，DynamicLossScaler会再次增加损失缩放比例。通过这种方式，DynamicLossScaler试图“骑在边缘”，始终使用尽可能高的损失缩放比例而不会导致溢出。

参数

• init_scale (float, 可选, default=2**32) – DynamicLossScaler尝试的初始损失缩放比例。

• scale_factor (float, 可选, default=2.0) – 调整损失缩放比例时使用的因子。如果遇到溢出，损失缩放比例会调整为损失缩放比例/scale_factor。如果在scale_window次连续迭代中没有发生溢出，损失缩放比例会调整为损失缩放比例*``scale_factor``。

• scale_window (int, 可选, default=1000) – 在增加损失缩放比例之前等待的没有溢出的连续迭代次数。

手动管理主参数

apex.fp16_utils.prep_param_lists(model, flat_master=False)

为给定模型创建FP32主参数列表，如Training Neural Networks with Mixed Precision: Real Examples中所述。

参数

• model (torch.nn.Module) – 现有Pytorch模型

• flat_master (bool, 可选, default=False) – 将主参数展平为一个张量，作为性能优化手段。

返回

返回一个元组 (model_params, master_params)。model_params是模型参数列表，供后续与model_grads_to_master_grads()和master_params_to_model_params()一起使用。master_params是FP32主梯度列表。如果flat_master=True，master_params将是包含一个元素的列表。

示例

model_params, master_params = prep_param_lists(model)

警告

目前，如果flat_master=True，模型的所有参数必须是相同类型。如果模型包含不同类型的参数，请使用flat_master=False，或使用FP16_Optimizer。

apex.fp16_utils.master_params_to_model_params(model_params, master_params, flat_master=False)

将主参数复制到模型参数。

参数

• model_params – 由prep_param_lists()创建的模型参数列表。

• master_params – 由prep_param_lists()创建的FP32主参数列表。如果master_params是使用flat_master=True创建的，那么也应将flat_master=True提供给master_params_to_model_params()。

apex.fp16_utils.model_grads_to_master_grads(model_params, master_params, flat_master=False)

将模型梯度复制到主梯度。

参数

• model_params – 由prep_param_lists()创建的模型参数列表。

• master_params – 由prep_param_lists()创建的FP32主参数列表。如果master_params是使用flat_master=True创建的，那么也应将flat_master=True提供给model_grads_to_master_grads()。