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

时间:2019-09-28 20:48:21

标签: python parallel-processing multiprocessing pytorch fast-ai

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

  

在单个CPU上运行代码(不执行多处理)仅需40秒即可处理将近50个   图片

     

使用割炬多处理功能在多个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()

我不确定是什么原因导致多处理模式下的速度下降。我读过很多讨论类似问题的文章,但在任何地方都找不到合适的解决方案。

1 个答案:

答案 0 :(得分:1)

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

意味着对于每个单个映像,您都将从磁盘中重新加载模型。 根据模型对象的大小,IO将花费更多的时间,然后再执行前进的步骤。

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