多进程包 - torch.multiprocessing¶
torch.multiprocessing 是对原生 multiprocessing
模块的包装器。
它注册了自定义的化简器,这些化简器使用共享内存来提供对不同进程中相同数据的共享视图。一旦张量/存储移动到共享内存(参见 share_memory_()
),就可以将其发送到其他进程而无需进行任何复制。
该 API 与原始模块 100% 兼容 - 只需将 import multiprocessing
更改为 import torch.multiprocessing
,即可通过队列或其他机制发送所有张量,并将它们移动到共享内存中。
由于 API 的相似性,我们没有记录此包的大部分内容,建议参考原始模块的优秀文档。
警告
如果主进程突然退出(例如,由于传入的信号),Python 的 multiprocessing
有时无法清理其子进程。这是一个已知的缺陷,因此,如果您在中断解释器后看到任何资源泄漏,则可能意味着这种情况已经发生在您身上。
策略管理¶
- torch.multiprocessing.set_sharing_strategy(new_strategy)[source]¶
设置 CPU 张量共享策略。
- 参数
new_strategy (str) – 所选策略的名称。应为
get_all_sharing_strategies()
返回的值之一。
共享 CUDA 张量¶
在进程之间共享 CUDA 张量仅在 Python 3 中受支持,使用 spawn
或 forkserver
启动方法。
与 CPU 张量不同,发送进程需要保留原始张量,只要接收进程保留张量的副本。引用计数是在后台实现的,但要求用户遵循以下最佳实践。
警告
如果消费者进程因致命信号而异常死亡,则只要发送进程正在运行,共享张量可能会永远保留在内存中。
尽快在消费者中释放内存。
## Good
x = queue.get()
# do somethings with x
del x
## Bad
x = queue.get()
# do somethings with x
# do everything else (producer have to keep x in memory)
2. 保持生产者进程运行,直到所有消费者退出。这将防止生产者进程释放消费者仍在使用的内存的情况。
## producer
# send tensors, do something
event.wait()
## consumer
# receive tensors and use them
event.set()
不要传递接收到的张量。
# not going to work
x = queue.get()
queue_2.put(x)
# you need to create a process-local copy
x = queue.get()
x_clone = x.clone()
queue_2.put(x_clone)
# putting and getting from the same queue in the same process will likely end up with segfault
queue.put(tensor)
x = queue.get()
共享策略¶
本节简要概述了不同共享策略的工作原理。请注意,它仅适用于 CPU 张量 - CUDA 张量将始终使用 CUDA API,因为这是它们唯一可以共享的方式。
文件描述符 - file_descriptor
¶
注意
这是默认策略(macOS 和 OS X 除外,它们不支持)。
此策略将使用文件描述符作为共享内存句柄。每当存储移动到共享内存时,从 shm_open
获取的文件描述符都会与对象一起缓存,并且当它要发送到其他进程时,文件描述符将传输(例如,通过 UNIX 套接字)到该进程。接收方也将缓存文件描述符并 mmap
它,以获得对存储数据的共享视图。
请注意,如果要共享大量张量,此策略将在大多数情况下保持大量文件描述符打开。如果您的系统对打开的文件描述符数量有较低的限制,并且您无法提高它们,则应使用 file_system
策略。
文件系统 - file_system
¶
此策略将使用传递给shm_open
的文件名来识别共享内存区域。这样做的好处是不需要实现缓存从中获得的文件描述符,但同时容易导致共享内存泄漏。文件不能在创建后立即删除,因为其他进程需要访问它来打开它们的视图。如果进程发生致命崩溃或被杀死,并且没有调用存储析构函数,则文件将保留在系统中。这非常严重,因为它们会一直占用内存,直到系统重新启动或手动释放它们。
为了解决共享内存文件泄漏的问题,torch.multiprocessing
将生成一个名为torch_shm_manager
的守护进程,该进程将与当前进程组隔离,并跟踪所有共享内存分配。一旦所有连接到它的进程退出,它将等待一段时间以确保没有新的连接,然后迭代该组分配的所有共享内存文件。如果发现其中任何一个仍然存在,则将释放它们。我们已经测试了这种方法,并证明它对各种故障具有鲁棒性。但是,如果您的系统具有足够高的限制,并且file_descriptor
是受支持的策略,我们建议不要切换到此策略。
生成子进程¶
注意
适用于 Python >= 3.4。
这取决于 Python 的multiprocessing
包中的spawn
启动方法。
可以通过创建Process
实例并调用join
来等待其完成,从而生成多个子进程来执行某些函数。这种方法在处理单个子进程时效果很好,但在处理多个进程时会出现潜在问题。
也就是说,顺序连接进程意味着它们将顺序终止。如果它们没有顺序终止,并且第一个进程没有终止,则进程终止将不会被注意到。此外,没有用于错误传播的原生工具。
下面的spawn
函数解决了这些问题,并处理错误传播、无序终止,并在检测到其中一个进程中的错误时主动终止进程。
- torch.multiprocessing.spawn.spawn(fn, args=(), nprocs=1, join=True, daemon=False, start_method='spawn')[source]¶
生成
nprocs
个进程,这些进程使用args
运行fn
。如果其中一个进程以非零退出状态退出,则其余进程将被杀死,并引发一个异常,其中包含终止原因。如果在子进程中捕获到异常,则会转发该异常,并且其回溯包含在父进程中引发的异常中。
- 参数
fn (function) –
函数被调用作为生成进程的入口点。此函数必须在模块的顶层定义,以便可以将其腌制并生成。这是多处理强加的要求。
函数被调用为
fn(i, *args)
,其中i
是进程索引,args
是传递的参数元组。args (tuple) – 传递给
fn
的参数。nprocs (int) – 要生成的进程数。
join (bool) – 对所有进程执行阻塞连接。
daemon (bool) – 生成的进程的守护进程标志。如果设置为 True,则将创建守护进程。
start_method (str) – (已弃用)此方法将始终使用
spawn
作为启动方法。要使用其他启动方法,请使用start_processes()
。
- 返回值
如果
join
为True
,则返回 None;如果join
为False
,则返回ProcessContext