如何提升Scipy稀疏数组的计算速度？附示例代码

我正在使用Scipy稀疏矩阵计算估计中用到的选择概率，当前执行10次耗时约15秒，请问如何优化以提升计算速度？

变量与参数 #

import itertools
import numpy as np
import scipy as sp
import timeit

T = 70000
N = 200
p = 212

X = sp.sparse.random(T, (N)*p, density=0.018, format='csr')
B = np.random.rand(p,)
I = sp.sparse.eye(N,format="csr")

原选择概率计算函数 #

def csr_example(X,B,I,tr):
    tem0 = sp.sparse.kron(I, B, format = "csr").transpose()
    exp_xb = X@tem0
    exp_xb.data = np.exp( exp_xb.data, out=exp_xb.data)
    exp_xb0 = sp.sparse.csr_matrix.sum(exp_xb,axis=1)
    exp_xb0 = np.reshape(exp_xb0, (tr, 1))
    pro = exp_xb/exp_xb0
    return pro

原执行耗时测试 #

print(timeit.timeit(lambda: csr_example(X,B,I,T), setup="pass",number=10))

执行10次耗时约15秒。

优化方案 #

核心优化点 #

1. 移除冗余的kron矩阵操作：原代码中kron(I, B)会创建规模为N*p × N的稀疏矩阵，该矩阵的乘法逻辑完全可以通过对X的列块直接与B做点积替代，避免了大矩阵的创建与额外乘法开销。
2. 改用稠密数组处理核心计算：T×N的结果数组（1400万元素）内存占用仅约112MB，完全可以加载到内存中，numpy的向量化操作比稀疏矩阵操作效率高一个数量级。
3. 原地操作减少内存开销：使用numpy原地计算避免额外内存分配，进一步提升速度。

优化后的代码 #

def optimized_csr_example(X, B, N, T, p):
    # 预分配结果数组，存储每个样本每个选择项的Xβ值
    xb = np.zeros((T, N), dtype=np.float64)
    # 遍历每个选择项，计算对应特征与B的点积
    for i in range(N):
        start_col = i * p
        end_col = start_col + p
        # 稀疏矩阵与稠密向量点积，结果直接存入数组
        xb[:, i] = X[:, start_col:end_col].dot(B).ravel()
    # 原地计算指数，节省内存
    np.exp(xb, out=xb)
    # 计算每行的和（概率分母），保持二维形状方便广播除法
    exp_sum = xb.sum(axis=1, keepdims=True)
    # 计算选择概率
    pro = xb / exp_sum
    # 若后续需要稀疏矩阵格式，可转换后返回；否则直接返回numpy数组
    return sp.sparse.csr_matrix(pro)

优化后的耗时测试 #

print(timeit.timeit(lambda: optimized_csr_example(X, B, N, T, p), setup="pass", number=10))

额外加速建议 #

• 使用优化的BLAS库：安装MKL版本的numpy/scipy（如通过conda安装mkl-service），矩阵点积操作会自动利用多线程加速。
• 避免不必要的稀疏转换：如果后续计算不需要稀疏矩阵格式，直接返回numpy数组即可，省去转换开销。

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