PyTorch 1.0 中文文档:多进程最佳实践

2022-05-07 14:11:43 浏览数 (1)

译者:cvley

torch.multiprocessing 是 Python 的 multiprocessing 的直接替代模块。它支持完全相同的操作,但进行了扩展,这样所有的张量就可以通过一个 multiprocessing.Queue 进行传递,将数据移动到共享内存并只将句柄传递到另一个进程。

注意

当一个 Tensor 传递到另一个进程时,Tensor 的数据是共享的。如果 torch.Tensor.grad 不是 None, 也会被共享。在一个没有 torch.Tensor.grad 域的 Tensor 被送到其他进程时,一个标准的进程专用的 .grad Tensor 会被创建,而它在所有的进程中不会自动被共享,与 Tensor 数据的共享方式不同。

这就允许实现各种训练方法,比如 Hogwild、A3C,或者其他那些需要异步操作的方法。

共享 CUDA 张量

进程间共享 CUDA 张量仅支持 Python 3,使用的是 spawn 或者 forkserver 启动方法。Python 2 中的 multiprocessing 仅使用 fork 来创建子进程,而 CUDA 运行时不支持该方法。

警告

CUDA API 需要分配给其他进程的显存在它们还在使用的情况下一直有效。你需要仔细确保共享的 CUDA 张量若非必须,不会超出使用范围。这对于共享模型参数不会是一个问题,但传递其他类型的数据时需要谨慎。注意该限制并不适用于共享 CPU 内存。

也可以参考:使用 nn.DataParallel 替代 multiprocessing

最佳实践和提示

避免和处理死锁

当创建一个新进程时,很多情况会发生,最常见的就是后台线程间的死锁。如果任何一个线程有锁的状态或者引入了一个模块,然后调用了fork,子进程很有可能处于中断状态,并以另外的方式死锁或者失败。注意即使你没这么做,Python 内建的库也有可能这么做——无需舍近求远,multiprocessing即是如此。multiprocessing.Queue 实际上是一个非常复杂的类,可以创建多个线程用于串行、发送和接收对象,它们也会出现前面提到的问题。如果你发现自己遇到了这种情况,尝试使用 multiprocessing.queues.SimpleQueue,它不会使用额外的线程。

我们在尽最大努力为你化繁为简,确保不会发生死锁的情况,但有时也会出现失控的情况。如果你遇到任何暂时无法解决的问题,可以在论坛上求助,我们将会研究是否可以修复。

阅读全文/改进本文

0 人点赞