使用Python multiprocessing.pool加速Pandas DataFrame列间线性回归

一、先修复多进程代码的错误 #

1. 必须添加if __name__ == '__main__'保护 #

Windows系统下使用multiprocessing时，必须将创建进程池的代码放在这个判断块内，避免重复初始化进程导致报错。

2. 修正lin_reg函数的参数问题 #

原函数参数命名混乱，且依赖全局变量col_y，多进程中全局变量传递不可靠，需要改为显式传参：

import pandas as pd
from scipy import stats
import multiprocessing

def lin_reg(col_x, col_y):
    lm = stats.linregress(col_x, col_y)
    return [
        col_x.name,
        lm.slope,
        lm.intercept,
        lm.rvalue,
        lm.pvalue,
        lm.stderr,
        lm.intercept_stderr
    ]

if __name__ == '__main__':
    # 假设这里是你的df1初始化代码
    # df1 = pd.read_csv(...)
    
    col_y = df1.iloc[:, 0]  # 第一列作为y
    # 准备所有任务的参数：每个任务是(x列, y列)的元组
    tasks = [(df1[col], col_y) for col in df1.columns[1:]]
    
    # 创建进程池
    with multiprocessing.Pool(4) as pool:
        # starmap会自动将每个元组的元素作为参数传给lin_reg
        results = pool.starmap(lin_reg, tasks)
    
    # 将结果转为DataFrame
    coef_df = pd.DataFrame(
        results,
        columns=["col", "slope", "intercept", "rvalue", "pvalue", "slope_stderr", "intercept_stderr"]
    )

错误原因说明 #

• 你之前的代码在循环里调用starmap，这完全误解了starmap的用法：starmap是一次性处理所有任务的可迭代序列，而非逐次调用。
• 未加if __name__ == '__main__'导致Windows系统下重复启动进程，触发RuntimeError。

二、更高效的向量化替代方案 #

多进程存在进程间数据拷贝的开销，对于大型DataFrame，向量化计算通常速度更快，CPU利用率更高。可以用numpy直接实现批量线性回归：

import pandas as pd
import numpy as np
from scipy import stats

def batch_linregress(x, y):
    # x是(n_samples, n_features)的数组，y是(n_samples,)的数组
    n = x.shape[0]
    y_mean = y.mean()
    x_means = x.mean(axis=0)
    
    # 计算协方差和方差
    cov = (x - x_means).T @ (y - y_mean) / (n - 1)
    var_x = np.var(x, axis=0, ddof=1)
    
    # 斜率和截距
    slope = cov / var_x
    intercept = y_mean - slope * x_means
    
    # 计算r值、p值、标准误
    r = cov / (np.sqrt(var_x) * np.std(y, ddof=1))
    df = n - 2
    t = r * np.sqrt(df / (1 - r**2))
    p = 2 * (1 - stats.t.cdf(np.abs(t), df=df))
    
    se_slope = np.sqrt((1 - r**2) / (df * var_x)) * np.std(y, ddof=1)
    se_intercept = se_slope * np.sqrt((x**2).mean(axis=0))
    
    return pd.DataFrame({
        "col": x.columns,
        "slope": slope,
        "intercept": intercept,
        "rvalue": r,
        "pvalue": p,
        "slope_stderr": se_slope,
        "intercept_stderr": se_intercept
    })

# 使用示例
col_y = df1.iloc[:, 0]
x_data = df1.iloc[:, 1:]
coef_df = batch_linregress(x_data, col_y)

这个方法直接用numpy的矩阵运算批量处理所有列，避免了循环和多进程的开销，速度会比多进程快很多，CPU利用率也能拉满。

三、其他优化建议 #

• 避免在循环中使用df.loc[i]添加行：原代码的coef_df.loc[i] = new_row会不断修改DataFrame，效率极低，应该先收集所有结果再一次性创建DataFrame（两种方案都采用了这种方式）。
• 如果DataFrame存在缺失值，先处理缺失值（比如df1.dropna()），否则stats.linregress会返回NaN，影响结果。

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