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