使用 pytorch 和多处理在 CPU 上运行推理

Question

我已经使用 FastAI（PyTorch 后端）在 GPU 上训练了一个 CNN 模型。我现在尝试在同一台机器上使用该模型进行推理，但使用 CPU 而不是 GPU。除此之外，我还尝试使用多处理模块来利用多个 CPU 内核。现在问题来了，

在单 CPU 上运行代码（无多处理）只需 40 秒即可处理近 50 张图像

使用 Torch 多处理在多个 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()

我不确定是什么导致了多处理模式下的速度变慢。我已经阅读了很多讨论类似问题的帖子，但在任何地方都找不到合适的解决方案。

Answer 1

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

这意味着对于每个图像，您都将从磁盘重新加载模型。根据模型对象的大小，IO 会比运行前向步骤更耗时。

请考虑在主线程中读取一次模型，然后使该对象可用于并行函数中的推理。

Answer 2

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

torch.set_num_threads(1)