译者: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
,它不会使用额外的线程。
我们在尽最大努力为你化繁为简,确保不会发生死锁的情况,但有时也会出现失控的情况。如果你遇到任何暂时无法解决的问题,可以在论坛上求助,我们将会研究是否可以修复。
阅读全文/改进本文