使用PyTorch和多进程在CPU上进行推理

Question

我使用FastAI（PyTorch后端）在GPU上训练了一个CNN模型。现在，我正在尝试在同一台机器上使用CPU而不是GPU对该模型进行推断。同时，我还尝试使用多个CPU核心，使用multiprocessing模块。现在遇到的问题是：

在单个CPU上运行代码（未使用multiprocessing）只需40秒即可处理近50张图像

使用torch multiprocessing在多个CPU上运行代码需要超过6分钟才能处理相同的50张图像

from torch.multiprocessing import Pool, set_start_method
os.environ['CUDA_VISIBLE_DEVICES']=""
from fastai.vision import *
from fastai.text import *
defaults.device = torch.device('cpu')

def process_image_batch(batch):

    learn_cnn  = load_learner(scripts_folder, 'cnn_model.pkl')
    learn_cnn.model.training = False    
    learn_cnn.model = learn_cnn.model.eval()
    # for image in batch: 
    #     prediction = ... # predicting the image here
    #     return prediction

if __name__ == '__main__':
    #
    # image_batches = ..... # retrieving the image batches (It is a list of 5 lists)
    # n_processes = 5
    set_start_method('spawn', force=True)
    try:
        pool = Pool(n_processes)
        pool.map(process_image_batch, image_batches)
    except Exception as e:
        print('Main Pool Error: ', e)
    except KeyboardInterrupt:
        exit()
    finally:
        pool.terminate()
        pool.join()

我不确定是什么原因导致多进程模式变慢。我已经阅读了很多帖子讨论类似的问题，但无法在任何地方找到合适的解决方案。

- asanoop24

---

你解决了吗？我也遇到了同样的问题。而且找不到合适的解决方案。我发现 PyTorch 在每个进程中都会尝试使用所有可用的 CPU，这可能导致巨大的减速，但我不确定... - AnarKi

是的，我必须强制PyTorch每个进程只使用一个线程。torch.set_num_threads(1) - asanoop24

仅使用一个线程在我的类似情况下也帮了我很大的忙。感谢您提供的信息。您可能希望将其发布（并接受）为自己问题的答案。 - Matthias

Answer 1

我认为你在这里犯了一个非常幼稚的错误，你正在函数中读取模型对象并对其进行并行处理。

这意味着对于每张图片，你都需要重新从磁盘中加载模型。根据模型对象的大小，I/O操作会比运行前向推理更耗时。

请考虑在主线程中仅读取一次模型，然后将该对象用于并行推理。

- tejas

---

@tejas...是的，那是我的错。我实际上是在使用批处理而不是单个图像。只是想在这里保持代码简单，但现在已经更新了。因此，模型是针对每个批次加载而不是每个图像，但仍然显示出糟糕的结果。 - asanoop24

所以你可以做一件简单的事情，假设你有8个核心和1600张图像需要推断。你要做的是将数据分成8个相等的部分，即每个部分200个文件。现在编写一个函数来加载模型对象，并在这200个文件上运行推断。最后使用多进程创建8个工作进程，并在1600个文件的8个块上并行化该函数。这样你只需要在每个进程中加载模型8次。 - tejas

Answer 2

解决方案是强制PyTorch每个进程仅使用1个线程，如下所示。

torch.set_num_threads(1)

- asanoop24