Scipy中Bool型CSR稀疏矩阵X@A@X计算内存不足问题求助

解决办法 #

1. 分块矩阵乘法（最直接有效） #

将大矩阵拆分为小行块，逐块计算(X_block @ A) @ X，最后拼接结果。这种方法避免一次性预分配超大内存，每块计算仅占用少量内存，适配你的设备配置。

示例代码：

import scipy.sparse as sparse

n = X.shape[0]
block_size = 100  # 可根据设备内存调整，建议在100-1000之间
result_blocks = []

for start in range(0, n, block_size):
    end = min(start + block_size, n)
    # 提取X的行块
    X_block = X[start:end, :]
    # 计算中间结果：X_block @ A
    temp = X_block @ A
    # 计算块结果：temp @ X
    block_result = temp @ X
    result_blocks.append(block_result)

# 拼接所有行块得到最终结果
result = sparse.vstack(result_blocks)

2. 利用对称性优化（进一步减少内存/计算量） #

由于X、A均为对称矩阵，X@A@X必然也是对称矩阵。只需计算上三角部分，再转置合并下三角，可减少约50%的计算量和内存占用。

示例代码：

import scipy.sparse as sparse
import numpy as np

n = X.shape[0]
block_size = 100
row = []
col = []
data = []

for start in range(0, n, block_size):
    end = min(start + block_size, n)
    X_block = X[start:end, :]
    temp = X_block @ A
    block_result = temp @ X
    
    # 转换为COO格式方便提取上三角
    coo_block = block_result.tocoo()
    # 修正行索引为原矩阵的行号
    original_row = coo_block.row + start
    # 筛选上三角元素（行 <= 列）
    upper_mask = original_row <= coo_block.col
    
    # 收集有效数据
    row.extend(original_row[upper_mask])
    col.extend(coo_block.col[upper_mask])
    data.extend(coo_block.data[upper_mask])

# 构建上三角COO矩阵
upper_tri = sparse.coo_matrix((data, (row, col)), shape=(n, n))
# 转置得到下三角，合并后去重
result = upper_tri + upper_tri.T
result.sum_duplicates()
# 转换为CSR格式方便后续操作
result = result.tocsr()

3. 手动计算点积（针对极小的A矩阵） #

由于A仅含2k个非零元素，可手动遍历A的非零位置，计算每个位置对结果的贡献并累加。此方法仅适用于A非零数极少的场景，否则内存占用会激增。

示例代码：

import scipy.sparse as sparse
import numpy as np

n = X.shape[0]
row = []
col = []
data = []

# 获取A的非零位置（COO格式）
A_coo = A.tocoo()
for i, j, val in zip(A_coo.row, A_coo.col, A_coo.data):
    # 获取X第i行的非零列索引（即X第i列的非零行索引，因X对称）
    rows_i = X.indices[X.indptr[i]:X.indptr[i+1]]
    # 获取X第j行的非零列索引
    rows_j = X.indices[X.indptr[j]:X.indptr[j+1]]
    
    # 生成外积的行和列索引
    r, c = np.meshgrid(rows_i, rows_j, indexing='ij')
    row.extend(r.ravel())
    col.extend(c.ravel())
    data.extend(np.full(len(r.ravel()), val, dtype=np.int64))

# 构建COO矩阵并合并重复位置的计数
result = sparse.coo_matrix((data, (row, col)), shape=(n, n))
result.sum_duplicates()
result = result.tocsr()

关键原理说明 #

Scipy稀疏矩阵乘法默认会保守预估结果的非零元素数量，当两个大稀疏矩阵相乘时，预估的非零数远大于实际值，导致预分配超大内存触发错误。分块计算通过将大矩阵拆分为小矩阵，每次仅处理少量数据，避免了一次性预分配超大内存；利用对称性则进一步减少了需要处理的数据量。

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