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

Question

为了充分利用 CPU/GPU，我运行了多个对不同数据集进行 DNN 推理（前馈）的进程。由于进程在前馈期间分配了 CUDA 内存，因此我收到了 CUDA 内存不足错误。为了缓解这种情况，我添加了torch.cuda.empty_cache()电话，这让事情变得更好。但是，仍然偶尔会出现内存不足错误。可能是由于错误的分配/发布时间。

我设法通过multiprocessing.BoundedSemaphore在前馈调用周围添加一个来解决这个问题，但这给初始化和进程之间共享信号量带来了困难。

在运行多个 GPU 推理过程时，是否有更好的方法来避免此类错误？

Answer 1

从我并行训练和推理的经验来看，想要榨干最后一点 GPU 内存几乎是不可能的。您能做的最好的事情可能就是估计可以并行运行的最大进程数，然后限制您的代码同时运行最多多个进程。使用信号量是限制并行进程数量并在有空槽时自动启动新进程的典型方法。

为了更容易地在进程之间初始化和共享信号量，您可以使用 amultiprocessing.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。然而，这可能会带来显着的性能开销，所以可能不是一个好主意。