多进程最佳实践

多进程最佳实践

CUDA 张量的共享
最佳实践和技巧

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()

译者署名

用户名	头像	职能	签名
风中劲草		翻译	人生总要追求点什么