ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

稀疏矩阵CSR格式从原理到实战:三数组压缩与图算法应用

稀疏矩阵CSR格式从原理到实战:三数组压缩与图算法应用 前几天在梳理图算法和推荐系统的底层实现时又一头扎进了稀疏矩阵的存储格式里。说实话稀疏矩阵这个名字听起来很学术但本质上就是一个非常朴素的工程问题当矩阵里大部分元素都是0时怎么用最少的内存和最快的速度把它存下来、算起来。在各种存储方案里CSRCompressed Sparse Row压缩稀疏行格式是出场率最高、也最值得先搞懂的一种。这篇文章我就把CSR格式掰开揉碎讲清楚它为什么是这样设计的、三个核心数组到底怎么配合、实际代码怎么写以及我在使用中踩过哪些坑。这篇内容适合正在学数据结构、需要处理大规模稀疏数据或者在写图算法、推荐系统、数值计算相关代码的同学参考。不管你是刚接触稀疏矩阵还是已经用过COO想进一步理解CSR这篇文章都能给你一套可以直接照做的思路和代码。1. 为什么稀疏矩阵值得单独学习在处理现实数据的时候稀疏矩阵几乎是躲不开的。用户和商品的关系、论文和关键词的关联、社交平台里人与人之间的关注这些场景的数据量动辄上亿但真正有值的位置往往少得可怜。如果老老实实用一个普通二维数组去存内存直接爆掉计算效率也会低得离谱。1.1 一个5阶矩阵的真实冲击先看一个最简单的例子。假设有一个5行5列的矩阵里面只有12个非零元素[ 1 0 0 0 2 ] [ 0 3 0 4 0 ] [ 5 0 6 0 0 ] [ 0 7 0 8 0 ] [ 9 0 0 0 10 ]用普通二维数组存储需要25个存储单元。如果矩阵规模变成100万行乘100万列还是只有100万个非零元素普通二维数组需要10的12次方个存储单元也就是大约1万亿个。即使每个元素只占4字节也需要4TB内存这个数据规模对绝大多数机器来说都是灾难。而CSR格式只存储非零值以及非零值的位置信息内存占用可以压缩到“非零元素数量级”这才是它真正的价值。1.2 零元素的存储浪费远超想象很多人觉得零元素不就是一个数字嘛存一下也没什么大不了。但问题的关键在于对于数字矩阵我们不仅需要存储值本身还需要存储它的行列坐标。如果做成坐标表每个非零元素都要存三维信息代价依然不低。CSR这一类压缩格式想解决的问题就是尽量用最紧凑的方式表达“哪些位置有值、值是什么、在哪一行”。CSR之所以叫Compressed Sparse Row本质就是“按行压缩稀疏矩阵”。它把矩阵看成一行一行的结构每行只记录本行里非零元素的“列号”和“值”最终用三个一维数组表达整个矩阵。理解了这个设计哲学后面所有的细节都会顺理成章。2. CSR格式的核心原理三数组协作CSR格式的精髓在于三个数组data也叫values、col_index也叫indices、row_ptr也叫indptr。这三个数组名在不同框架里有不同叫法SciPy里是data、indices、indptr很多C库里叫values、columns、row_pointer但含义完全一致。2.1 三个数组各自负责什么data把所有非零元素按行优先的顺序排成一维数组长度等于非零元素总数。col_index记录每个非零元素所在的列号长度和data相同data[i]在矩阵中的列位置就是col_index[i]。row_ptr记录每行第一个非零元素在data中的起始偏移位置长度是矩阵行数1。row_ptr[i]表示第i行的第一个非零元素在data中的下标row_ptr[i1] - row_ptr[i]就是第i行非零元素的个数。这里最反直觉的是row_ptr的设计它为什么不直接记录行数而要多存一个结束位置原因在于有了row_ptr[i]和row_ptr[i1]我们不需要额外的数组就能求出每一行的非零元素个数遍历一行时也可以通过这两个边界值确定循环范围。这是一种非常经典的“前闭后开区间”设计思路和C里begin()与end()的迭代器风格一致。2.2 手工推演一个3×4矩阵为了让这套逻辑落地我用手工推演一个3行4列的矩阵[ 1 0 0 2 ] [ 0 3 0 4 ] [ 5 0 6 0 ]按行优先扫描非零元素顺序为1, 2, 3, 4, 5, 6。那么data [1, 2, 3, 4, 5, 6]第0行的非零元素是1住第0列和2住第3列所以col_index前两个值是[0, 3]第1行的非零元素是3住第1列和4住第3列所以col_index接着是[1, 3]第2行的非零元素是5住第0列和6住第2列所以col_index接着是[0, 2]最终得到col_index [0, 3, 1, 3, 0, 2]row_ptr的第0个元素是0表示整个矩阵第一个非零元素从data[0]开始。第0行有2个非零元素区间是data[0]到data[1]所以row_ptr[1] 2第1行也是2个非零元素区间是data[2]到data[3]所以row_ptr[2] 4第2行也是2个非零元素区间是data[4]到data[5]所以row_ptr[3] 6最终row_ptr [0, 2, 4, 6]有了这三个数组恢复原矩阵的过程就是先看row_ptr确定每行边界再遍历这一行内部的非零元素用col_index确定列位置用data确定值。这个过程非常快因为不会遇到任何零元素。2.3 CSR和COO的区别与取舍在学CSR之前很多库默认先使用COOCoordinate List坐标列表格式。COO用三个等长数组分别存行号、列号、值比如(row, col, value)三元组。它的理解难度极低添加新元素也非常方便但缺点是空间浪费比较大因为每个元素都要存一个行号行号本身要占至少4字节100万行稀疏矩阵即使只存100万个非零元素光行号就要4MB加上列号还要4MB。CSR则通过row_ptr把“行号”压缩成了一种分段表示。原本每个非零元素都要记录行号现在只需要每个行边界记录一次。对于一个100万行、每行几十个非零元素的矩阵行号的存储开销从“非零元素数量×4字节”缩减为“行数×4字节”压缩效果非常明显。代价就是构建和修改的难度上升CSR里往矩阵中插入一个新元素通常需要移动大量数据所以它更适合一次性构建然后反复读取计算的场景而不是频繁动态修改的场景。从工程角度看常见做法是用COO在构建阶段不断往里添加数据构建完成后转成CSR格式进行高性能计算。SciPy的coo_matrix.tocsr()、PyTorch的torch.sparse_coo_tensor再到to_sparse_csr底层走的都是类似转换逻辑。3. 实操实现从COO构建CSR并完成转置3.1 先把COO转成CSR排序与压缩构建CSR最直接的方式是从一组“行号、列号、值”的三元组出发。假设我们有原始数据行号数组rows [0, 0, 1, 2, 2]列号数组cols [2, 0, 1, 0, 1]值数组vals [5, 1, 3, 6, 7]第一步永远是排序按行号为主键、列号为次键排序让同一行的元素连续排列并且每行内列号有序。排序之后三个数组变成rows [0, 0, 1, 2, 2]cols [0, 2, 1, 0, 1]vals [1, 5, 3, 6, 7]排完序后统计每一行的非零元素个数然后计算前缀和就能得到row_ptr。这里我建议你用Python做一次完整的模拟import numpy as np rows np.array([0, 0, 1, 2, 2], dtypenp.int32) cols np.array([2, 0, 1, 0, 1], dtypenp.int32) vals np.array([5, 1, 3, 6, 7], dtypenp.float32) # 按行和列排序 order np.lexsort((cols, rows)) rows_sorted rows[order] cols_sorted cols[order] vals_sorted vals[order] # 统计每行元素个数 row_count np.bincount(rows_sorted, minlength3) row_ptr np.zeros(len(row_count) 1, dtypenp.int32) np.cumsum(row_count, outrow_ptr[1:]) print(row_ptr:, row_ptr) print(col_index:, cols_sorted) print(data:, vals_sorted)输出结果是row_ptr: [0 2 3 5] col_index: [0 2 1 0 1] data: [1. 5. 3. 6. 7.]这里row_ptr[0]0表示第0行从data[0]开始row_ptr[1]2表示第0行有2个元素也就是data[0]和data[1]分别对应col_index[0]0和col_index[1]2。如果你把data[0]1放到矩阵第0行第0列data[1]5放到第0行第2列正好和原矩阵吻合。3.2 转置操作的实现思路矩阵转置就是把A[i][j]变成A_T[j][i]。放在CSR里最自然的思路是遍历原矩阵的每一个非零元素把它当作新矩阵中“列号原行号、行号原列号”的元素。换句话说转置后的新矩阵依然按行优先存储但新行号来自原矩阵的列号新列号来自原矩阵的行号。一个渐进实现方法是先遍历col_index统计每一列有多少个非零元素也就是新矩阵每一行有几个元素然后生成新的row_ptr。再次遍历所有非零元素把它们按新行号放入对应的位置同时把原行号写入新的col_index。这个过程听起来像是一个桶排序实际上就是计数排序在稀疏矩阵上的应用。如果矩阵比较小也可以直接用坐标转换再合并重复元素的方式构造(原列号, 原行号, 值)的三元组按新行号排序再重新计算row_ptr。我自己在写这种转置时会特别小心重复元素问题。如果原矩阵在同一个(row, col)位置出现多个值转置之后这些值也会落在同一个(new_row, new_col)位置这时候需要决定是相加、取最大值还是直接报错完全取决于业务诉求。SciPy默认是合并相加这一点在使用时要格外注意。3.3 性能实测CSR和普通数组的对比我自己在本地用1000万×1000万、非零元素数500万的随机稀疏矩阵做过一次标准内存对比测试。普通二维数组如果按双精度浮点存储理论内存约80TB直接不可行。CSR存储三个数组如果索引用4字节整型、值用8字节双精度内存占用大约是5,000,000 × (8 4) 10,000,001 × 4结果约为100MB左右。这组数据解释了为什么所有主流线性代数库和图计算引擎都选择CSR或CSC作为落地存储格式。如果你的矩阵规模很小比如只有几百行几百列CSR的构建开销反而可能超过直接存二维数组。这里需要给一个明确的结论当非零元素占比低于10%或者说矩阵规模大到二维数组存不下时CSR才有明显收益。更小的矩阵用普通数组代码还更简单。4. CSR在稀疏矩阵乘法与图算法中的实战CSR格式最大的优势不只是省内存而是让核心计算过程天然跳过零元素。稀疏矩阵和稠密向量相乘是PageRank、推荐系统、有限元分析里最常见的操作理解这个计算过程能让你对CSR的价值体验得特别真切。4.1 稀疏矩阵乘向量核心逻辑拆解假设矩阵A用CSR存储向量为x我们希望计算y A * x。普通矩阵乘法会对所有i, j做y[i] A[i][j] * x[j]其中大量A[i][j]等于0白白浪费时间。现代芯片乘法计算能力很强但访存带宽却是瓶颈跳过零元素能直接减少访存次数速度提升可以非常显著。用CSR实现乘向量的代码逻辑是外层循环遍历行通过row_ptr[i]和row_ptr[i1]确定当前行的元素区间内层循环遍历区间内的非零元素将data[k]与x[col_index[k]]相乘累加。写成C风格伪代码是for (int i 0; i n_rows; i) { y[i] 0; for (int k row_ptr[i]; k row_ptr[i1]; k) { y[i] data[k] * x[col_index[k]]; } }这个循环结构短时间内就可以写出来但性能差异体现在很多细节上比如col_index[k]对向量x的访问是不是连续这决定了缓存命中率再比如内层循环里是否有可以提前判断的零值分支。CSR之所以比COO适合乘法是因为COO格式里同一行的元素可能分裂在不同内存区域CPU必须跳来跳去而CSR保证了同一行的元素在内存里是紧密连续的。4.2 PageRank与图算法的关联很多学图计算的人会突然遇到CSR原因就在于图可以用邻接矩阵表示而稀疏图天然就是一个巨大的稀疏矩阵。一个千万级节点的社交网络邻接关系用普通二维矩阵根本无法存储但在CSR格式中row_ptr直接表示每个节点的出边起始位置col_index表示每条边指向的目标节点data可以是边的权重。PageRank算法每一次迭代的核心操作就是PR d * M * PR (1-d)/N其中M是列归一化后的转移矩阵。这一步完全就是稀疏矩阵乘向量。之前我调过一套PageRank代码用CSR格式实现在千万级边数据上迭代几十次只需要几十秒而如果用邻接表加逐条边遍历的方式即使逻辑一样因为缓存不友好整体耗时会翻几倍。这个差异会让第一次接触的人很有感触CSR从数据结构层面直接成就了工业级图算法。4.3 在推荐系统与数值计算中的延伸推荐系统里协同过滤算法会构建“用户-物品”稀疏矩阵行是用户列是物品值是评分或行为次数。计算用户相似度时不需要检查所有用户对所有物品的评分只需要对两个用户各自的非零物品集合做交集。如果有CSR格式每个用户的非零物品集合其实就是一个连续的数组片段配合二分查找或哈希可以非常高效地完成集合交集计算。CF里最常见的TopN相似用户计算底层用的就是这种思路。数值计算领域求解大型线性方程组时使用的有限元矩阵、稀疏迭代求解器里的预处理矩阵标准格式也基本都是CSR或类似变体。很多科学计算库比如Eigen、PETSc都提供了从COO到CSR的转换接口可见这套结构的普适性。5. 使用CSR过程中的常见坑与优化建议CSR看起来结构简单但真到实操环节问题往往出在那些容易被忽略的边界条件和索引语义上。我把自己踩过的坑整理了一下也提供一些优化思路方便大家少走弯路。5.1 千万不要忘记row_ptr的边界值新手最容易犯的错误是生成了row_ptr却忘记最后一位必须是非零元素总数。在遍历第i行时循环条件是k row_ptr[i1]如果row_ptr长度少了1或者最后一个值没有存总非零元素数程序就会越界访问。这个细节在数组越界时不容易排查因为C里可能没有立刻崩溃而读取到脏数据后会表现出非常诡异的数值错误。我自己习惯在构建CSR后立刻写一个自检函数逐行用row_ptr恢复稀疏矩阵然后和原始COO数据逐条比对位置和值。虽然多花一点点时间但能在早期发现问题。大数据工程师常说的“数据校验前置”在数据结构这里也一样成立。5.2 行内元素是否需要排序CSR格式本身不要求同一行内按列号排序但很多算法和库默认假设它是排好序的。比如在做矩阵乘法时部分优化库会提前检查col_index是否递增做二分查找时未排序的col_index会让结果完全错误。比如你写代码想判断A[i][j]是否等于某个值如果列号没排序线性扫描虽然也能找到答案但效率会很低一旦列号排好序就可以用upper_bound快速定位。具体到库的表现SciPy在tocsr()时会自动保证排序PyTorch的稀疏格式也会做规范化处理但如果你用底层C库尤其是自己拼接row_ptr和col_index时一定要检查排序情况。我在构建图算法时还有一个额外习惯按列号排序后col_index在GPU上的并行访问会更加规整避免线程发散。5.3 内存布局与缓存优化CSR在CPU上的性能瓶颈经常不是算法复杂度而是内存访问模式。矩阵乘向量时data[k]和col_index[k]是顺序访问的这一点对缓存很友好但x[col_index[k]]是随机访问向量x越大缓存命中率越低。当向量x无法完全放入L2缓存时性能会明显下降。一种常见优化是分块。把矩阵按行分成若干块每块对应的向量片段较小在遍历时反复使用能显著提升缓存命中率。另一种优化是调整存储布局把data改为16字节对齐把col_index压缩成int32甚至int16前提是矩阵规模和列号范围允许。GPU上还有一种常用技术就是给矩阵的行做重排序让相似行拥有相似的非零列集合提升并行线程的合并访存效率。这些优化不用一开始就做但如果你发现CSR乘法性能不理想就应该想到这个方向。5.4 动态修改CSR是灾难能避免就避免CSR一旦构建完成往里面增加一个新的非零元素就很困难因为可能需要在data数组中插入元素并移动后面所有的数据和row_ptr。如果业务需要频繁更新就别硬用CSR。一个更合理的方案是维护一个COO格式的“增量缓冲区”每次更新先写入缓冲区积累到一定量级后再统一merge到CSR里。很多工业级稀疏计算系统都是这么设计的它们会用日志式追加更新配合定期压缩思路和数据库里的LSM-Tree有异曲同工之处。如果你只是需要在已有非零元素上修改值那倒不用重建CSR因为data数组里对应位置可以直接赋值。但这里又有一个细节你必须先确认这个位置已经存在当前方案用哈希表维护三元组映射的话查询时间复杂度可以做到O(1)否则就得依靠row_ptr和col_index做逐行扫描。5.5 当心重复元素与非法索引数据构建过程中如果COO原始数据里包含重复的(row, col)转换到CSR时需要有明确的合并策略。数值计算里通常是求和图算法里可能是取第一条或累加权重。如果这个问题不提前定义好同一个输入在不同库里可能得到完全不同的矩阵这是数据一致性的大坑。另外原始数据里的row或col如果超出预期范围或者有负数都会导致row_ptr错乱。我在处理真实业务数据时会先在转换前做一轮索引清洗排除越界值和空值。这类问题看起来简单但一旦混进大规模数据里排查起来非常痛苦。6. 我常用的库与选择建议6.1 SciPy学习与原型验证首选Python生态里最常用的就是scipy.sparse。它提供csr_matrix、coo_matrix、csc_matrix等多个格式转换方便接口稳定。学习阶段我的建议是用coo_matrix构造原始数据再调用tocsr()进行转换。它内部会完成排序、去重、合并虽然我们不知道全部细节但结果永远是规范且正确的。from scipy.sparse import coo_matrix, csr_matrix row [0, 0, 1, 2, 2] col [2, 0, 1, 0, 1] data [5, 1, 3, 6, 7] A_coo coo_matrix((data, (row, col)), shape(3, 3)) A_csr A_coo.tocsr() print(data:, A_csr.data) print(indices:, A_csr.indices) print(indptr:, A_csr.indptr)输出data: [1. 5. 3. 6. 7.] indices: [0 2 1 0 1] indptr: [0 2 3 5]这就是一个最标准的CSR构造流程可以用来做各种验证实验。6.2 PyTorch和底层C库深度学习场景里PyTorch从1.9开始支持torch.sparse_csr_tensor可以直接构造CSR张量。在稀疏Transformer、图神经网络里这种格式被广泛使用。但要注意PyTorch的稀疏GPU算子覆盖度相比稠密算子还不够全面用之前最好确认你需要的操作有对应CUDA实现否则一调用就会落到CPU上性能反而更差。底层C库我常用Eigen和Intel MKL。Eigen的SparseMatrixdouble默认就是CSR使用setFromTriplets可以高效地从三元组构建。MKL提供了一整套稀疏BLAS接口比如mkl_dcsrmv做稀疏矩阵向量乘法性能非常接近理论峰值。如果你在做高性能数值计算这套组合是工业界的标准答案。7. 学习路径上的小技巧CSR格式之所以值得花时间不仅因为它本身常用更因为它背后代表了一套通用的压缩与索引思想。理解了CSR后再去学CSC压缩稀疏列、BSR分块稀疏行、DIA对角线都会快很多。CSC本质上是CSR的转置视角BSR则把一个小块当成一个元素来进一步压缩这些概念都是从三个数组的逻辑里生长出来的。我在学习时有一个非常有效的方法拿到任何一种稀疏格式首先手工写出它的“还原函数”也就是根据存储数组恢复原始二维矩阵。这一步能一次性验证你对格式每一个细节的理解。接下来再写“压缩函数”也就是从原始矩阵生成存储数组。两个函数来回验证格式就没有死角。最后再试着用CSR实现矩阵乘向量跟稠密矩阵的结果比较。这一套流程走下来你对CSR的理解会比看十篇文章都扎实。在实际项目中我还会提前评估数据规模和格式转换频率。如果只是做一次离线分析用SciPy一条龙解决就好如果要做在线服务就要考虑CSR构建成本以及如何避免重复构建。当你真正为千万级甚至亿级数据设计存储结构的时候CSR不只是一个知识点而是所有工程决策的起点。希望这篇文章能帮你把CSR这个基础但极其重要的工具吃透在之后的代码里真正用起来。
返回列表