SciPy稀疏矩阵csr_matrix与csc_matrix构造的中间内存占用问询

问题说明 #

使用SciPy提供的行、列索引构造方法创建csr_matrix或csc_matrix：

csr_matrix((data, (row_ind, col_ind)), [shape=(M, N)])
csc_matrix((data, (row_ind, col_ind)), [shape=(M, N)])

当前data+row_ind+col_ind合计占用25GB内存，剩余约35GB内存仍无法完成构造，调用函数时因内存溢出执行失败。官方文档未明确构造过程的中间内存消耗，需要分析具体的内存占用情况。

复现内存不足问题的示例代码：

import numpy as np
from scipy.sparse import csc_matrix

num_values = 2500000000
output_matrix_size = 150000

matrix = csc_matrix(
    (
        np.zeros(num_values, dtype=np.float16),
        (
            np.zeros(num_values, dtype=np.int32),
            np.zeros(num_values, dtype=np.int32),
        ),
    ),
    shape=(output_matrix_size, output_matrix_size),
)

内存占用分析 #

要理清构造过程的内存消耗，需结合CSR/CSC的最终存储结构和构造时的核心步骤：

1. 最终CSR/CSC的基础内存 #

• CSC格式：由三个核心数组构成：
  • data：存储非零元素值，长度等于非零元素数（即示例中的num_values）
  • indices：存储每个非零元素的行索引，长度等于num_values
  • indptr：存储每列的起始位置指针，长度为列数+1（示例中为150001）
• CSR格式：结构类似，indices存储列索引，indptr存储每行的起始位置指针，长度为行数+1

以示例代码为例，最终CSC矩阵的内存占用：

• data：25亿个float16，占25e9 * 2B = 50GB
• indices：25亿个int32，占25e9 *4B=100GB
• indptr：约586KB，可忽略不计
  最终总内存约150GB，但这只是最终结构的占用，构造过程的峰值内存更高。

2. 构造阶段的中间内存开销 #

构造稀疏矩阵时，SciPy的核心开销来自索引排序步骤：

• 排序临时数组：需要创建一个与索引数组长度一致的临时数组，用于存储排序后的索引（CSC对应行索引，CSR对应列索引）。示例中为25亿个int32，占25e9*4B=100GB。
• 重复索引处理（可选）：如果存在(row,col)重复的情况，还会额外占用内存存储聚合结果，但示例中全零索引会触发合并，不过核心开销仍来自排序。

加上原始输入的25GB，构造过程的峰值内存会达到原始25GB + 排序临时数组100GB + 最终结构的部分内存，远超过剩余的35GB，导致内存溢出。

优化建议 #

• 分块构造合并：将data、row_ind、col_ind拆分多个小块，分别构造小型CSR/CSC矩阵，再通过scipy.sparse.vstack/hstack合并，降低单步内存峰值。
• 提前预排序索引：预先对CSC的row_ind或CSR的col_ind排序，可让SciPy跳过内部排序步骤，大幅减少中间内存开销。
• 压缩索引数据类型：若矩阵维度允许，用int16替代int32存储索引，可将索引相关内存占用减半。

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