ProcessPoolExecutor比手写Process更稳,因其内置进程复用、超时控制、异常透传;避免AttributeError需将函数置于模块顶层且非lambda/嵌套;max_workers应据内存与任务粒度调优,而非盲目设为cpu_count()。
训练模型时 CPU 占用率“过高”不是 bug,而是 multiprocessing 正常工作的表现——它真正在压满多核,而不是像 threading 那样被 GIL 锁在单核上空转。关键问题是:高占用是否带来实际吞吐提升?有没有隐性代价(如内存爆炸、进程卡死)?下面直奔实操。
为什么 ProcessPoolExecutor 比手写 Process 更稳
你写 10 个 Process 启动模型推理,看似并行,但极易出问题:
- 每个
Process都要重新导入模块、加载模型、初始化环境,启动慢、内存翻倍(模型参数被 pickle 多次) - 没做异常捕获时,一个子进程崩溃会导致整个流程静默失败,
p.join()卡住不返回 - 无法自动限流:任务队列塞满、子进程堆积、OOM killer 开始杀进程
而 ProcessPoolExecutor 内置进程复用、超时控制、结果聚合和异常透传。它默认复用 worker 进程,模型只需在 initializer 里加载一次:
from concurrent.futures import ProcessPoolExecutor import osdef init_model(): global model model = load_heavy_ml_model() # 只执行一次/进程
def run_inference(data): return model.predict(data) # 复用已加载的 model
with ProcessPoolExecutor( max_workers=os.cpu_count() - 1, initializer=init_model ) as executor: results = list(executor.map(run_inference, data_batches))
Pool.map() 报 AttributeError: Can't pickle local object 怎么办
这是最常见报错,根本原因不是模型大,而是函数定义位置或形态不对:
- 函数不能定义在
if __name__ == '__main__':块内——子进程启动时会重新导入模块,找不到该函数 - 不能是 lambda、嵌套函数、类实例方法——它们无法被
pickle序列化 - 不能引用闭包变量(比如外层函数的局部变量),否则序列化失败
正确做法:把待并行函数放在模块顶层,参数显式传入:
# ✅ 正确:普通函数,参数全靠 args 传
def predict_on_batch(model_path, batch_data):
model = load_model(model_path) # 每次都加载?不推荐,见上条 initializer
return model(batch_data)
❌ 错误:在 main 里定义
if name == 'main':
def bad_func(x): return x * 2 # 子进程看不到
with Pool() as p:
p.map(bad_func, [1,2,3])
设多少个 max_workers 才不翻车
别直接填 os.cpu_count()。真实瓶颈常不在 CPU 核数,而在内存带宽或模型加载开销:
- 纯计算型(如数值仿真):用
os.cpu_count() - 1,留一核给系统调度和主进程响应 - 模型较大(>1GB)、GPU 显存有限时:从
2开始测,观察htop中 RES 列是否稳定、有无 OOM killer 日志(dmesg -T | grep -i "killed process") - 任务粒度极小(单次推理
Windows 下还必须加 if __name__ == '__main__':,否则子进程反复 fork 自身,触发 RecursionError 或直接崩溃。
模型加载后内存为何还是飙升
你以为子进程复用了模型,但若在 run_inference 里又调用了 torch.load() 或 tf.keras.models.load_model(),就等于每个任务都重复加载——内存必然线性增长。
真正节省内存的做法只有两个:
- 用
initializer在 worker 进程启动时加载一次,之后所有任务复用该全局变量(注意线程/进程安全) - 改用共享内存(如
multiprocessing.Array)存放只读权重,但需手动 flatten/reshape,工程成本高,仅适用于极特殊场景
另外提醒:子进程无法继承主进程的 CUDA 上下文,GPU 模型必须在每个子进程中独立 .to('cuda'),且显存也是按进程独占分配的——这点常被忽略,导致 GPU OOM 而非 CPU 占用高。