FastAI(PyTorch后端)CNN模型CPU多进程推理性能骤降问题排查

我来帮你分析下为什么多进程反而慢这么多，以及对应的解决办法：

核心问题分析 #

你的代码里有几个关键的性能瓶颈，直接导致了多进程效率远低于单进程：

• 每个进程重复加载模型：在process_image_batch函数里，每次进程处理一个batch时都会重新加载cnn_model.pkl模型。模型加载本身需要读取文件、反序列化权重、初始化计算图，这是非常耗时的操作——5个进程就会重复执行5次这个昂贵的步骤，光是加载模型的时间就占了大头。
• Spawn启动方式的额外开销：你用了set_start_method('spawn')，这种启动方式会让每个子进程重新初始化整个Python解释器环境，包括重新导入FastAI、PyTorch这些体积庞大的库，这又增加了大量的启动时间。
• CPU线程竞争冲突：PyTorch的CPU张量运算默认会通过OpenMP启用多线程加速，而如果你的多进程+每个进程的多线程同时运行，会导致CPU核心被过度占用，出现线程争抢资源的情况，反而拖慢了整体速度。

具体解决方案 #

1. 主进程一次性加载模型，子进程共享权重 #

避免每个进程重复加载模型，我们可以在主进程里先加载一次模型，然后把模型的权重传递给子进程，子进程只需要初始化模型并加载权重即可：

import os
import torch
from torch.multiprocessing import Pool, set_start_method
# 先设置线程数限制，避免后续冲突
os.environ['CUDA_VISIBLE_DEVICES'] = ""
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'

from fastai.vision import *

scripts_folder = "./your_script_dir"  # 替换成你的实际脚本目录
learn_cnn = None  # 子进程全局模型对象

def init_worker(model_weights):
    """子进程初始化：加载模型并导入权重"""
    global learn_cnn
    learn_cnn = load_learner(scripts_folder, 'cnn_model.pkl')
    learn_cnn.load_state_dict(model_weights)
    learn_cnn.model.eval()
    defaults.device = torch.device('cpu')

def process_image_batch(batch):
    predictions = []
    for image_path in batch:
        img = open_image(image_path)
        pred_class, pred_idx, outputs = learn_cnn.predict(img)
        predictions.append((image_path, pred_class))
    return predictions

if __name__ == '__main__':
    # 主进程先加载模型，获取权重
    learn_master = load_learner(scripts_folder, 'cnn_model.pkl')
    learn_master.model.eval()
    model_weights = learn_master.state_dict()

    set_start_method('spawn', force=True)
    image_batches = [["img1.jpg"], ["img2.jpg"], ...]  # 你的实际图片batch列表
    n_processes = 5

    try:
        # 用initializer传递模型权重给所有子进程，只初始化一次
        pool = Pool(n_processes, initializer=init_worker, initargs=(model_weights,))
        results = pool.map(process_image_batch, image_batches)
    except Exception as e:
        print('Main Pool Error: ', e)
    except KeyboardInterrupt:
        exit()
    finally:
        pool.terminate()
        pool.join()

2. 限制PyTorch的CPU线程数 #

在代码最开头设置OMP_NUM_THREADS和MKL_NUM_THREADS为1，确保每个子进程的模型推理只使用1个CPU线程，这样5个进程刚好占用5个核心，不会出现线程竞争的情况，最大化CPU利用率。

3. 可选：用多线程替代多进程 #

如果你的推理任务中图片IO（加载图片）占比很高，或者模型体积不大，多线程可能比多进程更高效——因为多线程不需要处理进程启动的额外开销，且FastAI模型在eval()模式下是线程安全的（没有反向传播，权重不会被修改）：

import os
os.environ['CUDA_VISIBLE_DEVICES'] = ""
from fastai.vision import *
from concurrent.futures import ThreadPoolExecutor

scripts_folder = "./your_script_dir"

def process_image(image_path):
    img = open_image(image_path)
    pred_class, pred_idx, outputs = learn_cnn.predict(img)
    return (image_path, pred_class)

if __name__ == '__main__':
    learn_cnn = load_learner(scripts_folder, 'cnn_model.pkl')
    learn_cnn.model.eval()
    defaults.device = torch.device('cpu')

    all_images = ["img1.jpg", "img2.jpg", ...]  # 所有图片路径列表
    with ThreadPoolExecutor(max_workers=5) as executor:
        results = list(executor.map(process_image, all_images))

效果验证 #

按照上面的修改后，你应该能看到多进程/多线程的耗时大幅降低，接近单进程耗时的1/5左右（理想状态）。核心就是解决了重复加载模型的开销和CPU资源竞争的问题。

内容的提问来源于stack exchange，提问作者asanoop24