PyTorch推理时多进程出现CUDA内存不足错误

Question

为了充分利用CPU/GPU，我运行了数个进程，在不同的数据集上进行DNN推理（前向传递）。由于这些进程在前向传递期间会分配CUDA内存，因此我遇到了CUDA内存溢出的错误。为了缓解这个问题，我添加了“torch.cuda.empty_cache()”函数调用来改善情况。然而，仍有偶尔出现内存不足的错误，可能是由于错误的内存分配/释放时间所致。

我通过在前向传递调用周围添加“multiprocessing.BoundedSemaphore”解决了这个问题，但这样会在进程之间初始化和共享信号量时引入难度。

在运行多个GPU推理进程时，有更好的方法可以避免这种错误吗？

- Xyand

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

另一方面，采用“贪婪”方法是在while循环中使用try ... except代码块并不断尝试使用GPU。然而，这可能会带来显著的性能开销，所以可能不是一个好主意。

- THN