分块结构稀疏矩阵外积累加计算的低内存优化咨询

最优低内存实现方案 #

核心思路是利用K远小于其他维度的特点，采用小批量分块计算的方式，既避免Python层的大量循环，又把中间矩阵的内存占用控制在极低水平。

方案优势 #

• 内存占用比现有向量化方案低2~3个数量级，中间最大矩阵仅为batch_size × K量级
• 运算速度和全向量化方案几乎持平，远快于原始逐块循环方案

实现代码 #

假设X为scipy.sparse.csr_matrix格式（按行切片效率最高，若为csc格式可先通过X.tocsr()转换）：

import numpy as np
from scipy.sparse import csr_matrix

# 批次大小可根据可用内存调整，因K极小，设为1024~10000都不会占用过多内存
batch_size = 2048
sum_Xi = np.zeros((K, K), dtype=X.dtype)

for batch_start in range(0, N, batch_size):
    batch_end = min(batch_start + batch_size, N)
    # 取出当前批次所有块对应的全部行
    batch_rows = X[batch_start * J : batch_end * J, :]
    # K极小的情况下转稠密开销可忽略，reshape后按块求和得到批次内各块的列和
    batch_arr = batch_rows.toarray()
    batch_s = batch_arr.reshape(-1, J, K).sum(axis=1)
    # 累加当前批次所有列和的外积之和
    sum_Xi += batch_s.T @ batch_s

特殊场景优化 #

如果X的稀疏度极高，转稠密反而带来额外开销，可以用以下版本，仅在批次内做少量循环：

batch_size = 2048
sum_Xi = np.zeros((K, K), dtype=X.dtype)

for batch_start in range(0, N, batch_size):
    batch_end = min(batch_start + batch_size, N)
    batch_n = batch_end - batch_start
    batch_s = np.zeros((batch_n, K), dtype=X.dtype)
    # 批次内仅循环batch_size次，远少于原始方案的N次循环
    for idx, i in enumerate(range(batch_start, batch_end)):
        batch_s[idx] = X[i*J : (i+1)*J, :].sum(axis=0).A1
    sum_Xi += batch_s.T @ batch_s

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