Python multiprocessing无法充分利用i9-13900K全部CPU核心

问题背景 #

使用Python的multiprocessing模块构建图数据集时，CPU核心无法充分利用：

• 硬件为i9-13900K（32线程），multiprocessing.cpu_count()返回32，系统通过lscpu -a --extended已识别32核心，sudo lshw | grep core显示configuration: cores=24 enabledcores=24 microcode=271 threads=32
• 实际仅2个核心处于运行状态，进程已创建但未占满CPU
• 尝试改用multiprocessing.Process()手动创建进程，问题依旧

相关代码 #

原始核心代码如下：

import mmcv
import os
from os import path as osp
import numpy as np
import torch
import torch.nn.functional as F
import pandas as pd
import TwoDimTTC
import warnings
import multiprocessing
warnings.filterwarnings("ignore")

def read_and_dump(idx, src_path, dst_path):
    data = mmcv.load(src_path + '/' + 'sample_infos_{}.pkl'.format(str(idx)))
    ego_velo = [data['can_bus'][14], data['can_bus'][13]]
    instance_info_list = data['instance_info']

    #### 构建图数据的代码，因篇幅省略 ####

    edge_index = torch.LongTensor(edge_index).t().contiguous()
    edge_type = torch.LongTensor(edge_type)
    edge_weight = torch.FloatTensor(edge_weight)

    graph_data_dict = {
        'node_features': node_features,
        'edge_type': edge_type,
        'edge_index': edge_index,
        'edge_weight': edge_weight,
    }
    mmcv.dump(graph_data_dict, osp.join(dst_path, 'sample_infos_{}.pkl'.format(str(idx))))

def create_graph(src_path, dst_path):
    pool = multiprocessing.Pool(processes=multiprocessing.cpu_count() - 1)
    for idx in range(len(os.listdir(src_path))):
        args = (idx, src_path, dst_path)
        pool.apply_async(func=read_and_dump, args=args)
    pool.close()
    pool.join()

if __name__ == '__main__':
    from tqdm import tqdm
    for d in tqdm(os.listdir('nuscenes_data_pkl_split/train/')):
        src = osp.join('nuscenes_data_pkl_split/train/', d)
        dst = osp.join('nuscenes_data_pkl_split/graph_split/train/', d)
        create_graph(src_path=src, dst_path=dst)

尝试改用multiprocessing.Process()的代码：

def create_graph(src_path, dst_path):
     ProcessWorker = []
     for idx in range(len(os.listdir(src_path))):
         p = multiprocessing.Process(target=read_and_dump, args=(idx, src_path, dst_path))
         ProcessWorker.append(p)
         p.start()
    
     for pw in ProcessWorker:
         pw.join()

解决方法 #

1. 排查任务负载与IO瓶颈 #

如果read_and_dump中的图构建逻辑计算量小，大部分时间消耗在mmcv.load/mmcv.dump等IO操作上，进程会频繁处于等待状态，无法占满CPU：

• 验证：在图构建部分添加CPU密集型测试代码（如大矩阵乘法），若测试时CPU能跑满，说明原逻辑存在IO瓶颈
• 优化：改用批量读写、使用更快的存储介质（如SSD），或把IO操作集中到主进程处理

2. 消除嵌套进程 #

当前代码外层循环遍历目录，每次调用create_graph创建新的进程池/多进程，形成进程嵌套，可能触发操作系统调度限制：

• 重构代码，统一收集所有任务后再创建进程池处理：

def main():
    from tqdm import tqdm
    tasks = []
    # 收集所有待处理任务
    for d in os.listdir('nuscenes_data_pkl_split/train/'):
        src = osp.join('nuscenes_data_pkl_split/train/', d)
        dst = osp.join('nuscenes_data_pkl_split/graph_split/train/', d)
        for idx in range(len(os.listdir(src))):
            tasks.append( (idx, src, dst) )
    
    # 统一创建进程池，预留2个核心给系统
    pool = multiprocessing.Pool(processes=multiprocessing.cpu_count() - 2)
    with tqdm(total=len(tasks)) as pbar:
        for _ in pool.imap_unordered(read_and_dump, tasks):
            pbar.update(1)
    pool.close()
    pool.join()

if __name__ == '__main__':
    main()

3. 限制第三方库的线程数 #

PyTorch、NumPy等第三方库默认会启用多线程，会与multiprocessing抢占CPU资源，导致核心利用率低：

• 在read_and_dump开头添加线程限制：

def read_and_dump(idx, src_path, dst_path):
    # 限制第三方库线程数，避免与多进程冲突
    torch.set_num_threads(1)
    os.environ['OMP_NUM_THREADS'] = '1'
    os.environ['MKL_NUM_THREADS'] = '1'
    
    # 原有代码...

4. 验证进程实际运行数量 #

在read_and_dump开头添加进程ID打印，确认实际运行的进程数是否符合预期：

def read_and_dump(idx, src_path, dst_path):
    print(f"Process {os.getpid()} handling idx {idx}")
    # 原有代码...

若打印的进程数远小于32，说明每个目录下的任务数量不足，导致进程池无法满负荷运行。

5. 检查CPU调度情况 #

i9-13900K包含大核与小核，操作系统可能优先调度进程到小核，导致表面利用率低：

• 使用htop工具查看每个核心的实时负载，确认是否真的只有2个核心在工作，还是显示问题。

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