python - 多处理时 Pytorch 推理 CUDA 内存不足

Question

为了充分利用 CPU/GPU，我运行了多个对不同数据集进行 DNN 推理(前馈)的进程。由于进程在前馈期间分配了 CUDA 内存，因此我收到了 CUDA 内存不足错误。为了缓解这种情况，我添加了 torch.cuda.empty_cache()电话让事情变得更好。但是，仍然偶尔会出现内存不足错误。可能是由于错误的分配/发布时间。
我设法通过添加 multiprocessing.BoundedSemaphore 解决了这个问题。围绕前馈调用，但这会给进程之间的初始化和共享信号量带来困难。
在运行多个 GPU 推理过程时，是否有更好的方法来避免此类错误？

Answer 1

从我并行训练和推理的经验来看，几乎不可能挤出最后一点 GPU 内存。可能你能做的最好的事情是估计可以并行运行的最大进程数，然后限制你的代码同时运行这么多进程。使用信号量是限制并行进程数量并在有空闲槽时自动启动新进程的典型方法。
为了更容易地在进程之间初始化和共享信号量，您可以使用 multiprocessing.Pool和池初始值设定项如下。

semaphore = mp.BoundedSemaphore(n_process)
with mp.Pool(n_process, initializer=pool_init, initargs=(semaphore,)) as pool:
    # here, each process can access the shared variable pool_semaphore

def pool_init(semaphore):
    global pool_semaphore
    pool_semaphore = semaphore

另一方面，贪婪的方法是使用 try ... except 运行。阻止 while循环并继续尝试使用 GPU。但是，这可能会带来显着的性能开销，因此可能不是一个好主意。