Pytorch多进程最佳实践

预备知识

模型并行（ model parallelism ）：即把模型拆分放到不同的设备进行训练，分布式系统中的不同机器（GPU/CPU等）负责网络模型的不同部分 —— 例如，神经网络模型的不同网络层被分配到不同的机器，或者同一层内部的不同参数被分配到不同机器，如AlexNet的训练。

数据并行（ data parallelism ）：即把数据切分，输入到不同的机器有同一个模型的多个副本，每个机器分配到不同的数据，然后将所有机器的计算结果按照某种方式合并。

多进程最佳实践

torch.multiprocessing 是 Python 的 multiprocessing 多进程模块的替代品。它支持完全相同的操作，但对其进行了扩展，以便所有通过多进程队列 multiprocessing.Queue 发送的张量都能将其数据移入共享内存，而且仅将其句柄发送到另一个进程。

注意：

当张量 Tensor 被发送到另一个进程时，张量的数据和梯度 torch.Tensor.grad 都将被共享。

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

一、CUDA 张量的共享

仅 Python 3 支持进程之间共享 CUDA 张量，我们可以使用 spawn 或forkserver 启动此类方法。 Python 2 中的 multiprocessing 多进程处理只能使用 fork 创建子进程，并且CUDA运行时不支持多进程处理。

警告：

CUDA API 规定输出到其他进程的共享张量，只要它们被这些进程使用时，都将持续保持有效。您应该小心并确保您共享的 CUDA 张量不会超出它应该的作用范围（不会出现作用范围延伸的问题）。这对于共享模型的参数应该不是问题，但应该小心地传递其他类型的数据。请注意，此限制不适用于共享的 CPU 内存。

也可以参阅： 使用 nn.DataParallel 替代多进程处理

二、最佳实践和技巧

1、避免和防止死锁

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

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

2、重用通过队列发送的缓冲区

请记住，每次将张量放入多进程队列 multiprocessing.Queue 时，它必须被移动到共享内存中。如果它已经被共享，将会是一个空操作，否则会产生一个额外的内存拷贝，这会减慢整个过程。即使您有一组进程将数据发送到单个进程，也可以让它将缓冲区发送回去，这几乎是不占资源的，并且可以在发送下一批时避免产生拷贝动作。

3、异步多进程训练（如： Hogwild）

使用多进程处理 torch.multiprocessing，可以异步地训练一个模型，参数既可以一直共享，也可以周期性同步。在第一种情况下，我们建议发送整个模型对象，而在后者中，我们建议只发送状态字典 state_dict() 。

我们建议使用多进程处理队列 multiprocessing.Queue 在进程之间传递各种 PyTorch 对象。使用 fork 启动一个方法时，它也可能会继承共享内存中的张量和存储空间，但这种方式也非常容易出错，应谨慎使用，最好只能让高阶用户使用。而队列，尽管它们有时候不太优雅，却能在任何情况下正常工作。

警告：

你应该留意没有用 if __name__ =='__main__' 来保护的全局语句。如果使用了不同于 fork 启动方法，它们将在所有子进程中执行。

4、Hogwild

具体的 Hogwild 实现可以在 示例库 中找到，但为了展示代码的整体结构，下面还有一个最简单的示例：

import torch.multiprocessing as mp
from model import MyModel

def train(model):
    # 构建 data_loader，优化器等
    for data, labels in data_loader:
        optimizer.zero_grad()
        loss_fn(model(data), labels).backward()
        optimizer.step()  # 更新共享的参数

if __name__ == '__main__':
    num_processes = 4
    model = MyModel()
    # 注意：这是 "fork" 方法工作所必需的
    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()

Reference

https://ptorch.com/news/176.html