PyTorch中用COO/CSR稀疏矩阵优化线性求解的内存节省方案

针对你当前用jacfwd生成稠密雅可比矩阵导致内存占用过高的问题，核心解决思路是仅计算并存储雅可比矩阵的非零元素，构建为COO/CSR稀疏矩阵后再执行稀疏线性求解，完全避免稠密矩阵的存储开销。以下是具体实现步骤：

1. 确定雅可比矩阵的非零结构 #

首先要明确雅可比矩阵中哪些位置是非零的：

• 如果你的残差函数get_residual有已知的稀疏模式（比如每个残差分量仅依赖少数未知量），可以直接手动定义非零元素的行索引rows和列索引cols（数组形状均为[N_nonzero]）。
• 如果不知道稀疏模式，可以先用小规模的unknown跑一次jacfwd，提取结果中的非零位置，后续大规模计算复用这个模式即可。

2. 计算稀疏雅可比矩阵的非零元素值 #

不用生成完整的稠密雅可比矩阵，而是针对非零位置批量计算偏导数：

方法一：用vmap批量计算单元素偏导 #

import torch
from functorch import vmap

unknown = unknown.requires_grad_(True)
r = get_residual(unknown)

# 定义单个位置的偏导计算函数
def get_jac_entry(row_idx, col_idx):
    return torch.autograd.grad(r[row_idx], unknown[col_idx], retain_graph=True)[0]

# 批量计算所有非零元素的值
jac_values = vmap(get_jac_entry)(rows, cols)

方法二：用jacrev结合索引提取 #

from functorch import jacrev

# 获取残差对未知量的雅可比梯度函数（不会显式生成完整矩阵）
jac_fn = jacrev(get_residual)
# 直接提取指定非零位置的元素值
jac_values = jac_fn(unknown)[rows, cols]

3. 构建稀疏矩阵并执行线性求解 #

PyTorch支持COO和CSR格式的稀疏矩阵，其中CSR格式在求解线性系统时效率更高：

构建COO/CSR稀疏矩阵 #

# 构建COO格式稀疏矩阵
J_coo = torch.sparse_coo_tensor(
    indices=torch.stack([rows, cols]),
    values=jac_values,
    size=(r.shape[0], unknown.shape[0]),
    device=unknown.device
)
# 转换为CSR格式（推荐用于求解）
J_csr = J_coo.to_sparse_csr()

稀疏线性系统求解 #

PyTorch的稠密求解器torch.linalg.solve不支持稀疏矩阵，需使用稀疏专用的迭代求解器：

# 使用GMRES求解 J * update = r（适合非对称矩阵）
update, _ = torch.sparse.linalg.gmres(J_csr, r, tol=1e-6)

# 或使用BiCGSTAB求解（适合对称正定矩阵）
update, _ = torch.sparse.linalg.bicgstab(J_csr, r, tol=1e-6)

关键优化细节 #

• 固定稀疏模式复用：若雅可比矩阵的非零结构在迭代过程中不变，仅需计算一次rows和cols，后续迭代只更新jac_values，大幅减少计算开销。
• 内存占用骤降：从稠密矩阵的O(n²)内存占用降到O(N_nonzero)，对于稀疏度1%以下的矩阵，内存节省可达两个数量级。
• 求解器选择：根据矩阵特性选迭代器，GMRES兼容性更强，BiCGSTAB在对称矩阵上收敛更快。

注意事项 #

• 批量计算偏导时，retain_graph=True可能导致内存泄漏，计算完成后可手动调用torch.cuda.empty_cache()（GPU场景）或清理无用张量。
• PyTorch稀疏矩阵的算子支持有限，若需复杂运算需提前验证兼容性。
• 若雅可比矩阵的非零结构动态变化，每次迭代都要重新计算rows和cols，需权衡稀疏方案的额外开销是否值得。

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