python 多进程独立启动和关闭进程

Question

我正在尝试使用多处理对 tensorflow 进行推理。 每个进程使用 1 个 GPU。 我有一个文件列表 input_files[]。 每个进程获取一个文件，在其上运行 model.predict 并将结果写入文件。 要继续下一个文件，我需要关闭进程并重新启动它。 这是因为 tensorflow 不允许 memory 的 go。 所以如果我使用相同的过程，我会得到 memory 泄漏。

我在下面写了一个有效的代码。 我启动 5 个进程，关闭它们并启动另一个 5。问题是所有进程都需要等待最慢的进程才能继续。 如何独立于其他进程启动和关闭每个进程？

请注意，Pool.map 超过 input_files_small 而不是 input_files。

file1 --> start new process --> run prediction --> close process --> file2 --> start new process --> etc.


for i in range(0, len(input_files), num_process):
    input_files_small = input_files[i:i+num_process]
    try:
        process_pool = multiprocessing.Pool(processes=num_process, initializer=init_worker, initargs=(gpu_ids))
        pool_output = process_pool.map(worker_fn, input_files_small)
    finally:
        process_pool.close()
        process_pool.join()

Answer 1

无需一遍又一遍地重新创建处理池。 首先，在创建池时指定maxtasksperchild=1 。 这应该会导致为每个提交的新任务创建一个新流程。 并且不要使用方法map ，而是使用不会阻塞的方法map_async 。 如果您的工作人员 function 没有返回您需要的结果，您可以使用pool.close后跟pool.join()来等待这些提交完成，如下所示或使用第二个代码变体：

process_pool = multiprocessing.Pool(processes=num_process, initializer=init_worker, initargs=(gpu_ids), maxtasksperchild=1)
for i in range(0, len(input_files), num_process):
    input_files_small = input_files[i:i+num_process]
    process_pool.map_async(worker_fn, input_files_small))
# wait for all outstanding tasks to complete
process_pool.close()
process_pool.join()

如果您需要worker_fn的返回值：

process_pool = multiprocessing.Pool(processes=num_process, initializer=init_worker, initargs=(gpu_ids), maxtasksperchild=1)
results = []
for i in range(0, len(input_files), num_process):
    input_files_small = input_files[i:i+num_process]
    results.append(process_pool.map_async(worker_fn, input_files_small))
# get return values from map_async
pool_outputs = [result.get() for result in results]
# you do not need process_pool.close() and process_pool.join()

但是，由于稍后调用map_async的map_async启动时，可能仍有一些“缓慢”的任务仍在运行，因此其中一些任务可能仍需要等待才能运行。 但至少池中的所有进程都应该保持相当繁忙。

如果您希望您的工作人员 function 出现异常并且需要在您的主进程中处理它们，那么它会变得更加复杂。