Python多进程训练神经网络时训练环节耗时异常增加的问题排查

我用Python multiprocessing做神经网络训练，子进程负责构造batch样本并放入队列，父进程从队列读取数据后用PyTorch训练。但发现多进程总耗时并没有比单进程短，进一步排查发现：多进程中从队列读取数据确实比单进程构造样本快，但训练环节（代码和单进程完全一致）的耗时反而更长。

复现代码 #

import multiprocessing as mp
import numpy as np
import time

n = 200

def get_sample():
    local_loop = 400
    # data
    x = np.random.rand(n,n)
    p = np.random.rand(n,n)
    y = 0
    for i in range(local_loop):
        y += np.power(x, p)
    return y

def new_process(q_data, total_loops):
    for i in range(total_loops):
        q_data.put(get_sample())
    print('finish new process')

def main(multi_proc=False):
    st = time.time()
    total_loops = 100
    
    local_loop = 2500
    mt = 0
    other_t = 0

    st_multi = time.time()
    if multi_proc:
        q_data = mp.Queue()
        new_proc = mp.Process(target=new_process,args=(q_data, total_loops))
        new_proc.start()
    mt += time.time() - st_multi

    for i in range(total_loops):
        st_multi = time.time()
        if multi_proc:
            y = q_data.get()
        else:
            y = get_sample()
        mt += time.time() - st_multi

        other_st = time.time()
        for j in range(local_loop):
            y += np.random.rand(n,n)
        other_t += time.time() - other_st

    st_multi = time.time()
    if multi_proc:
        assert q_data.empty()
        new_proc.join()
    mt += time.time() - st_multi

    print('\nmulti_proc', multi_proc)
    print('multi_proc_time', mt)
    print('other_time', other_t)

    print(f'total time: {time.time()-st}')

if __name__ == '__main__':
    main(multi_proc=False)
    main(multi_proc=True)

运行结果 #

multi_proc False
multi_proc_time 36.44150114059448
other_time 39.08155846595764
total time: 75.5232412815094
finish new process

multi_proc True
multi_proc_time 0.4313678741455078
other_time 40.54900646209717
total time: 40.980711460113525

可以看到，multi_proc=True时，模拟训练环节的other_time比单进程时长1秒以上，且该现象在多平台多次实验中一致，实际场景中已抵消多进程的收益。

原因分析 #

• 数据序列化与内存布局问题：mp.Queue通过pickle传递numpy数组，父进程拿到的是反序列化后的新数组，其内存布局（比如是否连续）可能和单进程直接生成的数组不同。numpy的运算对内存连续性很敏感，非连续数组会导致缓存命中率下降，运算速度变慢。
• CPU资源竞争：子进程构造样本时会占用CPU核心，父进程训练时需要和子进程抢占CPU资源，尤其是核心数较少的机器，上下文切换和资源争抢会降低训练环节的执行效率。
• numpy多线程冲突：numpy默认启用多线程加速运算，多进程模式下，子进程和父进程的numpy线程会同时抢占CPU核心，导致父进程的训练运算无法充分利用硬件资源，增加耗时。

验证与解决思路 #

1. 检查并修复内存布局：在父进程读取数据后，将数组转换为连续内存布局：

   if multi_proc:
       y = q_data.get()
       y = np.ascontiguousarray(y)  # 转换为连续数组
   

   重新运行后观察other_time是否下降。
2. 限制numpy线程数：在代码开头添加np.set_num_threads(1)，避免多进程间的numpy线程竞争，让每个进程只使用单线程运算，减少资源争抢。
3. 改用共享内存传递数据：使用multiprocessing.shared_memory或mp.Array来传递样本数据，避免pickle序列化/反序列化的开销和内存拷贝，父进程直接操作共享内存中的数据，保持内存布局一致。

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