非常游戏网
为什么Python训练模型时CPU占用率过高_利用Multiprocessing多进程并发

为什么Python训练模型时CPU占用率过高_利用Multiprocessing多进程并发

2026-05-29日常编程220263

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 os

def 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 占用高。