如何在SciPy中加速稀疏矩阵与稠密向量的子矩阵乘法？

问题根源 #

你的测试结果符合SciPy CSR矩阵的特性：

1. 列切片开销大：CSR格式是为行操作优化的，列切片需要重新构建整个稀疏矩阵的indptr、indices和data数组，这个预处理开销会抵消乘法的提速收益。
2. 小计算量下的Overhead：当d很小时，乘法本身的计算量极低，SciPy的dot函数的Python层开销占比极高，导致提速不明显。而Matlab的稀疏矩阵实现对这类场景有更细致的底层优化，因此能获得显著提速。

优化方案 #

1. 避免显式创建子矩阵，直接利用原矩阵结构计算

跳过创建sub_A的步骤，直接通过原CSR矩阵的索引和数据计算结果，彻底消除切片开销。

import scipy.sparse as sp
import numpy as np
import time

m = 10000
n = 100
d = 100
times = 100
x = np.ones((m,1))

A = sp.random(n, m, density=0.01, format='csr')

# 原方法的选中列和子向量
row_indices = np.random.choice(m, d, replace=False)
sub_x = x[row_indices].flatten()

# 优化方法1：全维度映射计算
start_time = time.time()
# 预处理：将选中列的x值映射到全维度数组（仅需一次）
x_full = np.zeros(m)
x_full[row_indices] = sub_x
for i in range(times):
    c = A.dot(x_full)
end_time = time.time()
print("直接映射计算耗时:", end_time - start_time)

# 优化方法2：无映射直接计算（避免创建全维度数组）
row_indices_sorted = np.sort(row_indices)
sub_x_sorted = sub_x[np.argsort(row_indices)]

start_time = time.time()
for i in range(times):
    c = np.zeros(n)
    # 遍历每一行，累加选中列的贡献
    for row in range(n):
        start_idx = A.indptr[row]
        end_idx = A.indptr[row+1]
        cols = A.indices[start_idx:end_idx]
        # 找到当前行中属于选中列的位置
        mask = np.isin(cols, row_indices_sorted)
        if np.any(mask):
            # 匹配sub_x中的对应值
            pos = np.searchsorted(row_indices_sorted, cols[mask])
            c[row] = np.sum(A.data[start_idx:end_idx][mask] * sub_x_sorted[pos])
end_time = time.time()
print("无映射直接计算耗时:", end_time - start_time)

2. 使用Numba加速自定义乘法逻辑

用Numba编译自定义的稀疏矩阵乘法函数，跳过SciPy的底层开销，针对你的场景做针对性优化。

import numba as nb

@nb.njit(parallel=True)
def csr_selected_dot(indptr, indices, data, selected_cols, x_selected, out):
    n_rows = len(indptr) - 1
    for row in nb.prange(n_rows):
        start = indptr[row]
        end = indptr[row+1]
        total = 0.0
        for idx in range(start, end):
            col = indices[idx]
            # 二分查找判断列是否在选中集合中
            low, high = 0, len(selected_cols)-1
            while low <= high:
                mid = (low + high) // 2
                if selected_cols[mid] == col:
                    total += data[idx] * x_selected[mid]
                    break
                elif selected_cols[mid] < col:
                    low = mid + 1
                else:
                    high = mid - 1
        out[row] = total

# 预处理排序（二分查找需要有序）
row_indices_sorted = np.sort(row_indices)
sub_x_sorted = sub_x[np.argsort(row_indices)]
out = np.zeros(n)

start_time = time.time()
for i in range(times):
    csr_selected_dot(A.indptr, A.indices, A.data, row_indices_sorted, sub_x_sorted, out)
end_time = time.time()
print("Numba加速耗时:", end_time - start_time)

3. 使用PyTorch的稀疏矩阵实现

深度学习框架的稀疏矩阵模块对切片和乘法有更高效的底层优化，适合这类场景：

import torch

# 转换为PyTorch CSR稀疏矩阵
A_torch = torch.sparse_csr_tensor(
    torch.from_numpy(A.indptr),
    torch.from_numpy(A.indices),
    torch.from_numpy(A.data),
    size=(n, m),
    dtype=torch.float32
)
sub_x_torch = torch.from_numpy(sub_x).float()

# 直接选取子矩阵并计算
sub_A_torch = A_torch[:, row_indices]

start_time = time.time()
for i in range(times):
    c = sub_A_torch @ sub_x_torch
end_time = time.time()
print("PyTorch加速耗时:", end_time - start_time)

4. GPU加速（CuPy）

如果有GPU可用，CuPy的稀疏矩阵操作能带来数量级的提速：

import cupy as cp
import cupyx.scipy.sparse as cps

# 转换为CuPy CSR矩阵
A_cp = cps.csr_matrix(A)
sub_x_cp = cp.array(sub_x)
sub_A_cp = A_cp[:, row_indices]

start_time = time.time()
for i in range(times):
    c = sub_A_cp.dot(sub_x_cp)
cp.cuda.Stream.null.synchronize()  # 等待GPU计算完成
end_time = time.time()
print("CuPy GPU加速耗时:", end_time - start_time)

效果预期 #

采用方案1或方案2，在d < m/100的场景下，速度提升能接近Matlab的水平；方案3和4则能进一步突破CPU性能瓶颈，获得更大的收益。

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