利用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将比运行前一步花费更多时间。

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