Python使用statsmodels.GMM实现因果中介分析时Jupyter崩溃问题

云原生阿岚

2026-10-6

问题描述

开展因果中介分析时需要使用GMM（广义矩估计）估算处理效应，编写代码如下：

from statsmodels.sandbox.regression.gmm import GMM
import numpy as np
import pandas as pd

class GMMAB(GMM):

    def __init__(self, *args, **kwds):
        # 设置矩条件和参数的对应数量
        kwds.setdefault('k_moms', 6)
        kwds.setdefault('k_params', 6)
        super(GMMAB, self).__init__(*args, **kwds)


    def momcond(self, params):
        c = params
        y,m = self.endog.T #[y,m]
        x = self.exog.squeeze() # x
        
        g1 = m - c[1] - c[0]*x
        g2 = x*(m - c[1] - c[0]*x)
        g3 = y - c[2] - c[3]*x - c[4]*m- c[5]*m*x
        g4 = x*(y - c[2] - c[3]*x - c[4]*m- c[5]*m*x)
        g5 = m*(y - c[2] - c[3]*x - c[4]*m- c[5]*m*x)
        g6 = m*x*(y - c[2] - c[3]*x - c[4]*m- c[5]*m*x)
        g = np.column_stack((g1, g2, g3, g4, g5, g6))
        return g

beta0 = np.array([0.1, 0.1, 0.1, 0.1, 0.1, 0.1])
dta = pd.read_csv('mediation_data.csv')
y = np.array(dta.y) # y,m,x shape: [100000, 1] 
m = np.array(dta.m)
s = np.array(dta[['y','m']])
x = np.array(dta.x)
model = GMMAB(endog = s, exog = x, instrument = x, k_moms=6, k_params=6)

beta0 = np.array([0.1,0.1,0.1,0.1,0.1,0.1])
model.fit(beta0, maxiter=2, weights_method='hac', optim_method='nm')

数据集包含10万条以上观测，y、m、x均为长度10万以上的数组，运行环境可用内存大于3G，在Jupyter Notebook中运行上述代码时未抛出任何异常，但Notebook直接崩溃，无法确定是GMM实现逻辑有误还是内存不足导致崩溃，以下是对应排查及解决建议：

排查及解决建议

小样本验证逻辑正确性：先从全量数据中随机抽取1000~10000条观测，替换原有数据集运行代码，如果可以正常得到结果，则排除代码逻辑错误问题，基本定位为大样本下的内存溢出导致崩溃。
优化矩条件函数的冗余计算：现有代码中momcond方法内重复计算了m和y的残差项，可将公共残差提前计算存储，减少重复计算开销：

def momcond(self, params):
    c = params
    y,m = self.endog.T
    x = self.exog.squeeze()
    # 提前计算公共残差
    resid_m = m - c[1] - c[0]*x
    resid_y = y - c[2] - c[3]*x - c[4]*m - c[5]*m*x
    g1 = resid_m
    g2 = x * resid_m
    g3 = resid_y
    g4 = x * resid_y
    g5 = m * resid_y
    g6 = m * x * resid_y
    return np.column_stack((g1, g2, g3, g4, g5, g6))

调整权重计算方法避免内存爆炸：你当前使用的weights_method='hac'是导致大样本内存溢出的最常见原因，statsmodels的sandbox版GMM实现HAC权重时，会生成大规模中间协方差矩阵，10万样本下很容易超出3G内存限制。可先将权重方法替换为weights_method='unadjusted'或weights_method='robust'验证运行流程，确实能正常运行后，如果必须使用HAC权重，可手动设置更小的bw_args带宽参数降低计算量，或换用性能优化更好的线性模型库的GMM接口。
替换更高效的求解器：当前使用的optim_method='nm'（Nelder-Mead单纯形法）属于无梯度优化，迭代次数多、内存开销大，可替换为optim_method='bfgs'等梯度类求解器，收敛速度更快，大幅降低迭代过程中的内存占用。
降低单次计算的内存开销：调试阶段可先降低maxiter参数值，验证流程无问题后再逐步提升迭代次数确保收敛。

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