# 多进程最佳实践

> 原文：[`pytorch.org/docs/stable/notes/multiprocessing.html`](https://pytorch.org/docs/stable/notes/multiprocessing.html)> 
>
> 译者：[飞龙](https://github.com/wizardforcel)
>
> 协议：[CC BY-NC-SA 4.0](http://creativecommons.org/licenses/by-nc-sa/4.0/)


`torch.multiprocessing`是 Python 的[`multiprocessing`](https://docs.python.org/3/library/multiprocessing.html#module-multiprocessing "(在 Python v3.12 中)")模块的一个替代品。它支持完全相同的操作，但扩展了它，使得通过[`multiprocessing.Queue`](https://docs.python.org/3/library/multiprocessing.html#multiprocessing.Queue "(在 Python v3.12 中)")发送的所有张量的数据都会移动到共享内存中，并且只会发送一个句柄给另一个进程。

注意

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

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

## 多进程中的 CUDA

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

注意

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

与 CPU 张量不同，发送进程需要保留原始张量，只要接收进程保留了张量的副本。这是在底层实现的，但需要用户遵循程序正确运行的最佳实践。例如，发送进程必须保持活动状态，只要消费者进程引用了张量，如果消费者进程通过致命信号异常退出，引用计数无法帮助您。请参阅此部分。

另请参阅：使用 nn.parallel.DistributedDataParallel 代替 multiprocessing 或 nn.DataParallel

## 最佳实践和提示

### 避免和解决死锁

当生成一个新进程时，有很多事情可能会出错，最常见的死锁原因是后台线程。如果有任何持有锁或导入模块的线程，并且调用了`fork`，那么子进程很可能处于损坏状态，并且会发生死锁或以不同方式失败。请注意，即使您没有，Python 内置库也会 - 不需要查看比[`multiprocessing`](https://docs.python.org/3/library/multiprocessing.html#module-multiprocessing "(在 Python v3.12 中)")更远的地方。[`multiprocessing.Queue`](https://docs.python.org/3/library/multiprocessing.html#multiprocessing.Queue "(在 Python v3.12 中)")实际上是一个非常复杂的类，它生成多个线程用于序列化、发送和接收对象，它们也可能引起上述问题。如果您发现自己处于这种情况，请尝试使用`SimpleQueue`，它不使用任何额外的线程。

我们正在尽力让您轻松使用，并确保这些死锁不会发生，但有些事情超出我们的控制。如果您遇到无法解决的问题，请尝试在论坛上寻求帮助，我们会看看是否可以解决。

### 通过队列传递的缓冲区

请记住，每次将`Tensor`放入[`multiprocessing.Queue`](https://docs.python.org/3/library/multiprocessing.html#multiprocessing.Queue "(在 Python v3.12 中)")时，它必须被移动到共享内存中。如果已经是共享的，则不会执行任何操作，否则将产生额外的内存复制，可能会减慢整个过程的速度。即使您有一组进程向单个进程发送数据，也要让它发送缓冲区回来 - 这几乎是免费的，并且可以避免在发送下一批数据时进行复制。

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

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

我们建议使用[`multiprocessing.Queue`](https://docs.python.org/3/library/multiprocessing.html#multiprocessing.Queue "(在 Python v3.12 中)")来在进程之间传递各种 PyTorch 对象。例如，当使用`fork`启动方法时，可以继承已经在共享内存中的张量和存储，但这很容易出错，应谨慎使用，仅供高级用户使用。队列，即使它们有时不够优雅，也会在所有情况下正常工作。

警告

您应该注意全局语句，如果没有用`if __name__ == '__main__'`保护，它们将在所有子进程中执行。

#### Hogwild

在[示例存储库](https://github.com/pytorch/examples/tree/master/mnist_hogwild)中可以找到一个具体的 Hogwild 实现，但为了展示代码的整体结构，下面也有一个最小示例：

```py
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 过度订阅是一个技术术语，指的是分配给系统的虚拟 CPU 总数超过硬件上可用的虚拟 CPU 总数的情况。

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

在[Hogwild 实现](https://github.com/pytorch/examples/tree/main/mnist_hogwild)中的代码示例中查看 CPU 过度订阅。

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

```py
python  main.py  --num-processes  4 
```

假设机器上有 N 个虚拟 CPU 可用，执行上述命令将生成 4 个子进程。每个子进程将为自己分配 N 个虚拟 CPU，导致需要 4*N 个虚拟 CPU。然而，机器上只有 N 个虚拟 CPU 可用。因此，不同的进程将竞争资源，导致频繁的进程切换。

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

1.  高 CPU 利用率：通过使用`htop`命令，您可以观察到 CPU 利用率始终很高，经常达到或超过其最大容量。这表明对 CPU 资源的需求超过了可用的物理核心，导致进程之间的争用和竞争。

1.  低系统效率下的频繁上下文切换：在 CPU 过度订阅的情况下，进程竞争 CPU 时间，操作系统需要快速在不同进程之间切换以公平分配资源。这种频繁的上下文切换会增加开销并降低整个系统的效率。

### 避免 CPU 过度订阅

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

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

假设机器上有 N 个虚拟 CPU，并且将生成 M 个进程，每个进程使用的最大`num_threads`值将为`floor(N/M)`。为了避免在 mnist_hogwild 示例中出现 CPU 过度订阅，需要对[示例存储库](https://github.com/pytorch/examples/tree/main/mnist_hogwild)中的`train.py`文件进行以下更改。

```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 替换为可用的虚拟 CPU 数量，将 M 替换为选择的进程数量。适当的`num_thread`值将根据手头的具体任务而变化。然而，作为一般准则，`num_thread`的最大值应为`floor(N/M)`，以避免 CPU 过度订阅。在[mnist_hogwild](https://github.com/pytorch/examples/tree/main/mnist_hogwild)训练示例中，在避免 CPU 过度订阅后，您可以实现 30 倍的性能提升。