多进程最佳实践¶

torch.multiprocessing 是 Python multiprocessing 模块的直接替换。它支持完全相同的操作，但对其进行了扩展，以便所有通过 multiprocessing.Queue 发送的张量，其数据都将移动到共享内存中，并且只会向另一个进程发送句柄。

注意

当一个 Tensor 被发送到另一个进程时，Tensor 数据是共享的。如果 torch.Tensor.grad 不是 None，它也会被共享。在一个没有 torch.Tensor.grad 字段的 Tensor 被发送到另一个进程后，它会创建一个标准的进程特定的 .grad Tensor，该张量不会像 Tensor 的数据那样自动在所有进程之间共享。

这允许实现各种训练方法，如 Hogwild、A3C 或任何其他需要异步操作的方法。

多进程中的 CUDA¶

CUDA 运行时不支持 fork 启动方法；需要使用 spawn 或 forkserver 启动方法才能在子进程中使用 CUDA。

注意

启动方法可以通过使用 multiprocessing.get_context(...) 创建上下文或直接使用 multiprocessing.set_start_method(...) 来设置。

与 CPU 张量不同，只要接收进程保留张量的副本，发送进程就需要保留原始张量。这在底层实现，但需要用户遵循最佳实践才能使程序正确运行。例如，只要消费者进程有对张量的引用，发送进程就必须保持活动状态，并且如果消费者进程通过致命信号异常退出，则引用计数无法拯救您。请参阅本节。

另请参阅：使用 nn.parallel.DistributedDataParallel 而不是 multiprocessing 或 nn.DataParallel

最佳实践和技巧¶

避免和对抗死锁¶

当生成新进程时，有很多事情可能会出错，最常见的死锁原因是后台线程。如果任何线程持有锁或导入模块，并且调用了 fork，则子进程很可能处于损坏状态，并将死锁或以其他方式失败。请注意，即使您不这样做，Python 内置库也会这样做 - 无需进一步查看 multiprocessing。multiprocessing.Queue 实际上是一个非常复杂的类，它生成多个线程用于序列化、发送和接收对象，它们也可能导致上述问题。如果您发现自己处于这种情况，请尝试使用 SimpleQueue，它不使用任何额外的线程。

我们正在尽力为您提供便利，并确保这些死锁不会发生，但有些事情是我们无法控制的。如果您有任何暂时无法解决的问题，请尝试在论坛上联系我们，我们将看看是否可以解决问题。

重用通过队列传递的缓冲区¶

请记住，每次将 Tensor 放入 multiprocessing.Queue 时，它都必须移动到共享内存中。如果它已经是共享的，则它是空操作，否则将产生额外的内存复制，这可能会减慢整个过程。即使您有一个进程池向单个进程发送数据，也要使其将缓冲区发送回去 - 这几乎是免费的，并且可以让您在发送下一批数据时避免复制。

异步多进程训练（例如 Hogwild）¶

使用 torch.multiprocessing，可以异步训练模型，参数可以一直共享，也可以定期同步。在第一种情况下，我们建议发送整个模型对象，而在后一种情况下，我们建议仅发送 state_dict()。

我们建议使用 multiprocessing.Queue 在进程之间传递各种 PyTorch 对象。例如，当使用 fork 启动方法时，可以继承已经存在于共享内存中的张量和存储，但这非常容易出错，应谨慎使用，并且仅供高级用户使用。队列，即使它们有时不是最优雅的解决方案，但在所有情况下都能正常工作。

警告

您应该注意是否有全局语句，这些语句没有用 if __name__ == '__main__' 保护。如果使用了 fork 以外的其他启动方法，它们将在所有子进程中执行。

Hogwild¶

一个具体的 Hogwild 实现可以在 examples 仓库中找到，但为了展示代码的整体结构，下面也提供了一个最小示例

import torch.multiprocessing as mp
from model import MyModel

def train(model):
    # Construct data_loader, optimizer, etc.
    for data, labels in data_loader:
        optimizer.zero_grad()
        loss_fn(model(data), labels).backward()
        optimizer.step()  # This will update the shared parameters

if __name__ == '__main__':
    num_processes = 4
    model = MyModel()
    # NOTE: this is required for the ``fork`` method to work
    model.share_memory()
    processes = []
    for rank in range(num_processes):
        p = mp.Process(target=train, args=(model,))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()

多进程中的 CPU¶

不适当的多进程可能导致 CPU 过度订阅，导致不同的进程争夺 CPU 资源，从而导致效率低下。

本教程将解释什么是 CPU 过度订阅以及如何避免它。

CPU 过度订阅¶

CPU 过度订阅是一个技术术语，指的是分配给系统的 vCPU 总数超过硬件上可用的 vCPU 总数的情况。

这会导致对 CPU 资源的严重争用。在这种情况下，进程之间频繁切换，这会增加进程切换开销并降低整体系统效率。

在示例仓库中找到的 Hogwild 实现的代码示例中查看 CPU 过度订阅。

当使用以下命令在 CPU 上使用 4 个进程运行训练示例时

python main.py --num-processes 4

假设机器上有 N 个 vCPU 可用，执行上述命令将生成 4 个子进程。每个子进程将为其自身分配 N 个 vCPU，从而导致需要 4*N 个 vCPU。但是，机器只有 N 个 vCPU 可用。因此，不同的进程将争夺资源，导致频繁的进程切换。

以下观察结果表明存在 CPU 过度订阅

高 CPU 利用率：通过使用 htop 命令，您可以观察到 CPU 利用率始终很高，通常达到或超过其最大容量。这表明对 CPU 资源的需求超过了可用的物理核心，从而导致进程之间争用和竞争 CPU 时间。
频繁的上下文切换和低系统效率：在 CPU 过度订阅的情况下，进程争夺 CPU 时间，操作系统需要快速在不同进程之间切换，以公平地分配资源。这种频繁的上下文切换增加了开销并降低了整体系统效率。

避免 CPU 过度订阅¶

避免 CPU 过度订阅的一个好方法是适当的资源分配。确保并发运行的进程或线程数不超过可用的 CPU 资源。

在这种情况下，一个解决方案是在子进程中指定适当的线程数。这可以通过在子进程中使用 torch.set_num_threads(int) 函数为每个进程设置线程数来实现。

假设机器上有 N 个 vCPU，并且将生成 M 个进程，则每个进程使用的最大 num_threads 值将为 floor(N/M)。为了避免 mnist_hogwild 示例中的 CPU 过度订阅，需要对 example 仓库中的文件 train.py 进行以下更改。

def train(rank, args, model, device, dataset, dataloader_kwargs):
    torch.manual_seed(args.seed + rank)

    #### define the num threads used in current sub-processes
    torch.set_num_threads(floor(N/M))

    train_loader = torch.utils.data.DataLoader(dataset, **dataloader_kwargs)

    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
    for epoch in range(1, args.epochs + 1):
        train_epoch(epoch, args, model, device, train_loader, optimizer)

使用 torch.set_num_threads(floor(N/M)) 为每个进程设置 num_thread。其中，您将 N 替换为可用的 vCPU 数量，将 M 替换为选择的进程数量。适当的 num_thread 值将根据手头的具体任务而有所不同。但是，作为一般准则，num_thread 的最大值应为 floor(N/M)，以避免 CPU 过度订阅。在 mnist_hogwild 训练示例中，在避免 CPU 过度订阅后，您可以实现 30 倍的性能提升。