10

我已经使用 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()

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

4

1 回答 1

5

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

对于每一个图像来说,这意味着您正在从磁盘重新加载模型。根据您的模型对象大小,IO 会比向前迈出更多时间。

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

于 2019-09-30T09:11:41.033 回答