优化空间统计模型中Python函数的内存占用问题

问题背景 #

实现BYM2模型缩放因子计算时，针对全美国207972×207972规模的邻接矩阵（稀疏度0.003%，含1314488个非零元素），遇到以下问题：

• 直接调用spsolve(Q_perturbed, b)求解全量单位矩阵时，即使分配750GB内存仍触发内存溢出；
• 改用逐列循环求解的方式仅需60GB内存，但速度极慢——20分钟仅完成5%计算量；
• 尝试CSR/CSC稀疏格式、降低数据类型等手段，未解决核心问题。

核心瓶颈分析 #

1. 全量求解的内存爆炸：求解n个右侧向量会生成稠密的n×n矩阵（207972²个float32元素约占160GB），且spsolve的中间计算过程需要数倍内存，远超集群分配上限。
2. 逐列求解的重复开销：每次调用spsolve都会重复执行矩阵分解（LU分解等），这是性能低下的核心原因——分解仅需执行一次，后续求解可复用结果。

优化方案 #

1. 复用矩阵分解结果 #

使用scipy.sparse.linalg.splu对Q_perturbed提前做LU分解，后续所有右侧向量求解直接复用分解后的对象，彻底消除重复分解的开销。

2. 避免存储全量伪逆矩阵 #

最终仅需Q_inv的对角线元素的几何均值，无需存储完整的Sigma（Q_perturbed的逆矩阵）。通过逐个求解单位向量，直接计算Q_inv的对角线元素，大幅降低内存占用。

3. 精细化内存管理 #

• 避免将稀疏矩阵转为稠密数组，直接使用稠密向量作为右侧输入；
• 及时删除不再使用的变量，结合gc.collect()主动回收内存。

优化后的代码实现 #

from scipy.sparse import diags, csc_matrix
from scipy.sparse.linalg import splu
import numpy as np
import gc
import logging

def scaling_factor_sp_optimized(A):
    """Compute the scaling factor from an adjacency matrix for BYM2 model.
    Only works for fully connected graphs.
    """
    # 构建稀疏精度矩阵
    num_neighbors = A.sum(axis=1).A.ravel().astype(np.float32)
    D = diags(num_neighbors, format="csc", dtype=np.float32)
    Q = D - A
    del num_neighbors, D, A
    gc.collect()

    # 添加对角扰动保证矩阵可逆
    jitter = max(Q.diagonal()) * np.sqrt(np.finfo(np.float32).eps)
    Q_perturbed = Q + diags(np.ones(Q.shape[0]) * jitter, dtype=np.float32, format="csc")
    del jitter, Q
    gc.collect()

    n = Q_perturbed.shape[0]
    # 执行LU分解，复用分解结果
    lu = splu(Q_perturbed.tocsc())
    del Q_perturbed
    gc.collect()

    # 计算W = Sigma @ 1，等价于求解Q_perturbed @ W = 1
    ones_vec = np.ones(n, dtype=np.float32)
    W = lu.solve(ones_vec)
    sum_W = W.sum()
    del ones_vec
    gc.collect()

    # 累计Q_inv对角线元素的对数和
    log_diag_sum = 0.0
    for i in range(n):
        if i % 10**4 == 0:
            logging.info(f"Processing row {i}")
            print(f"Processing row {i}")
        # 生成单位向量e_i
        e_i = np.zeros(n, dtype=np.float32)
        e_i[i] = 1.0
        # 求解Sigma的第i列
        sigma_i = lu.solve(e_i)
        # 计算Q_inv的对角线元素
        q_inv_diag = sigma_i[i] - (W[i] ** 2) / sum_W
        log_diag_sum += np.log(q_inv_diag)
        # 及时释放临时变量
        del e_i, sigma_i
        gc.collect()

    # 清理剩余资源
    del lu, W, sum_W
    gc.collect()

    # 计算几何均值并返回
    return np.exp(log_diag_sum / n)

scaling_factor = scaling_factor_sp_optimized(adj_matrix)

预期优化效果 #

• 内存占用：无需存储n×n的稠密矩阵，内存占用可控制在60GB以内；
• 求解速度：仅执行一次LU分解，后续单向量求解开销大幅降低，预计速度比原逐列方案提升10-100倍；
• 稳定性：避免全量求解的内存溢出问题，无需等待2TB内存节点。

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