apex.parallel

class apex.parallel.DistributedDataParallel(module, message_size=10000000, delay_allreduce=False, shared_param=None, allreduce_trigger_params=None, retain_allreduce_buffers=False, allreduce_always_fp32=False, num_allreduce_streams=1, allreduce_communicators=None, gradient_average=True, gradient_predivide_factor=1.0, gradient_average_split_factor=None, prof=False)[源代码]

apex.parallel.DistributedDataParallel是一个模块包装器，可以轻松实现多进程分布式数据并行训练，类似于torch.nn.parallel.DistributedDataParallel。参数在初始化时跨参与进程广播，并且梯度在backward()期间被 allreduce 并跨进程平均。

DistributedDataParallel针对与 NCCL 一起使用进行了优化。它通过在backward()期间将通信与计算重叠，并将较小的梯度传输分桶以减少所需的传输总数来实现高性能。

DistributedDataParallel旨在与上游启动实用程序脚本torch.distributed.launch一起使用，其中--nproc_per_node <= 每个节点的 gpu 数量。当与此启动器一起使用时，DistributedDataParallel假定进程与 GPU 的 1:1 映射。它还假定您的脚本在创建模型之前调用torch.cuda.set_device(args.rank)。

https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed显示了详细用法。https://github.com/NVIDIA/apex/tree/master/examples/imagenet显示了另一个结合了DistributedDataParallel与混合精度训练的示例。

参数

  • module – 要在多 GPU/分布式模式下运行的网络定义。

  • message_size (int, default=1e7) – 通信桶中的最小元素数。

  • delay_allreduce (bool, default=False) – 将所有通信延迟到向后传递结束。这会禁用通信与计算的重叠。

  • allreduce_trigger_params (list, optional, default=None) – 如果提供，应包含从模型中提取的参数列表。每当这些参数之一收到其梯度时（而不是当大小为 message_size 的存储桶已满时）将启动 Allreduce。在 backward() 的末尾，还将自动执行清理 allreduce 以捕获任何剩余的梯度。如果提供了 allreduce_trigger_params，则将忽略 message_size 参数。

  • allreduce_always_fp32 (bool, default=False) – 在 allreduce 之前将任何 FP16 梯度转换为 FP32。这可以提高广泛扩展运行的稳定性。

  • gradient_average (bool, default=True) – 用于切换 DDP 是否对进程的 allreduce 梯度进行平均的选项。对于正确的缩放，建议使用默认值 True。

  • gradient_predivide_factor (float, default=1.0) – 允许在 allreduce 之前和之后部分地对进程的梯度进行平均。在 allreduce 之前：grads.mul_(1.0/gradient_predivide_factor)。在 allreduce 之后：grads.mul_(gradient_predivide_factor/world size)。这可以减少广泛扩展运行的 FP16 allreduce 动态范围上的压力。

警告

如果gradient_average=False，则仍将应用 allreduce 前的除法（grads.mul_(1.0/gradient_predivide_factor)），但将省略 allreduce 后的梯度平均（grads.mul_(gradient_predivide_factor/world size)）。

forward(*inputs, **kwargs)[源代码]

定义每次调用时执行的计算。

应该由所有子类重写。

注意

尽管需要在该函数中定义正向传递的配方，但应在此之后调用Module实例，而不是调用此函数，因为前者负责运行注册的钩子，而后者会默默地忽略它们。

class apex.parallel.Reducer(module_or_grads_list)[源代码]

apex.parallel.Reducer是一个简单的类，可以帮助 allreduce 模块的参数跨进程。Reducer旨在为用户提供额外的控制：与DistributedDataParallel不同，Reducer不会在backward()期间自动 allreduce 参数。相反，Reducer等待用户手动调用<reducer_instance>.reduce()。例如，这使得可以将 allreduce 延迟到每几个迭代而不是每个迭代执行。

与DistributedDataParallel一样，Reducer平均其 allreduce 的任何张量，超过参与进程的数量。

Reducer旨在与上游启动实用程序脚本torch.distributed.launch一起使用，其中--nproc_per_node <= 每个节点的 gpu 数量。当与此启动器一起使用时，Reducer假定进程与 GPU 的 1:1 映射。它还假定您的脚本在创建模型之前调用torch.cuda.set_device(args.rank)。

参数

module_or_grads_list – 既可以是在多 GPU/分布式模式下运行的网络定义（module），也可以是要减少的梯度 iterable。如果传入的是 module，Reducer 构造函数会将参数同步到各个进程（从 rank 0 广播），以确保它们都使用相同的值进行初始化。如果传入的是梯度列表（来自某个 module），用户需要负责在训练开始时手动同步该 module 的参数。

class apex.parallel.SyncBatchNorm(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False, fuse_relu=False)[source]

同步批量归一化模块，从 torch.nn.BatchNormNd 扩展而来，增加了跨多个进程的统计量缩减。apex.parallel.SyncBatchNorm 旨在与 DistributedDataParallel 一起使用。

在训练模式下运行时，该层会减少所有进程的统计量，以增加归一化层的有效批量大小。这在给定进程上的批量大小较小，会降低模型收敛准确性的应用中很有用。该模型使用 torch.distributed 中的集合通信包。

在评估模式下运行时，该层会回退到 torch.nn.functional.batch_norm

参数

  • num_features – 来自预期大小为 \((N, C, L)\) 的输入的 \(C\)，或来自大小为 \((N, L)\) 的输入的 \(L\)

  • eps – 添加到分母以获得数值稳定性的值。默认值：1e-5

  • momentum – 用于 running_mean 和 running_var 计算的值。可以设置为 None 以进行累积移动平均（即简单平均）。默认值：0.1

  • affine – 布尔值，当设置为 True 时，此模块具有可学习的仿射参数。默认值：True

  • track_running_stats – 布尔值，当设置为 True 时，此模块跟踪运行均值和方差；当设置为 False 时，此模块不跟踪此类统计信息，并且始终在训练和评估模式下都使用批量统计信息。默认值：True

  • process_group – 传入一个进程组，在其中同步小批量的统计信息。None 表示使用默认进程组

  • channel_last – 布尔值，当设置为 True 时，此模块将输入张量的最后一个维度作为通道维度。默认值：False

示例：
>>> # channel first tensor
>>> sbn = apex.parallel.SyncBatchNorm(100).cuda()
>>> inp = torch.randn(10, 100, 14, 14).cuda()
>>> out = sbn(inp)
>>> inp = torch.randn(3, 100, 20).cuda()
>>> out = sbn(inp)
>>> # channel last tensor
>>> sbn = apex.parallel.SyncBatchNorm(100, channel_last=True).cuda()
>>> inp = torch.randn(10, 14, 14, 100).cuda()
forward(input, z=None)[source]

定义每次调用时执行的计算。

应该由所有子类重写。

注意

尽管需要在该函数中定义正向传递的配方，但应在此之后调用Module实例，而不是调用此函数，因为前者负责运行注册的钩子，而后者会默默地忽略它们。

实用函数

apex.parallel.convert_syncbn_model(module, process_group=None, channel_last=False)[source]

递归遍历模块及其子模块，以将所有 torch.nn.modules.batchnorm._BatchNorm 实例替换为 apex.parallel.SyncBatchNorm。

所有 torch.nn.BatchNorm*N*d 都封装在 torch.nn.modules.batchnorm._BatchNorm 周围，因此此函数使您可以轻松切换到使用同步 BN。

参数

module (torch.nn.Module) – 输入模块

示例

>>> # model is an instance of torch.nn.Module
>>> import apex
>>> sync_bn_model = apex.parallel.convert_syncbn_model(model)