大规模稀疏矩阵最小二乘计算加速方案咨询

背景 #

我通过最小二乘法求解方程A_des·x = b以得到x，该操作需执行数千次。每次迭代中A_des会变化，但始终保持稀疏性（矩阵中仅约4%为有效数据）。

问题 #

A_des规模较大（维度为(28106, 1185)），单次求解线性系统耗时约1秒，而我需要执行30万次该操作。

咨询问题 #

请问有哪些可行的加速方案？

已尝试方案 #

我分别在CPU和GPU上实现了求解函数，使用pytorch进行GPU计算，但加速效果不稳定。由于求解器通常已支持多线程，无法额外并行化。在CPU上尝试了多种求解器：np.linalg.solve、np.linalg.lstsq、scipy.linalg.solve、scipy.optimize.least_squares、scipy.sparse.linalg.lsqr（预先将A_des转换为scipy.sparse.coo_array(A_des)）。

当前最优方案 #

我找到的最优折中方案是使用np.linalg.solve，但需要手动进行矩阵操作。例如，用最小二乘函数求解为x = np.linalg.lstsq(A_des, b)，而用solve则需写成x = np.linalg.solve(A_des.T@A_des, A_des.T@b)。另一个更快的选项是x = scipy.sparse.linalg.lsqr(A_des, b, atol = 1e-3, btol = 1e-3)，但求解结果精度不佳。

我的求解函数 #

def Inverter(GPU, A_des, b):

    if GPU:
        # Migrate b to torch
        b= torch.from_numpy(b).to(device).double()
        # Migrate design matrix to GPU
        A_des = torch.from_numpy(A_des).to(device)
        x= torch.linalg.solve(A_des.T@A_des,A_des.T@(b)).cpu()
    else:
        #x= scipy.sparse.coo_array(A_des)
        #x= scipy.sparse.linalg.lsqr(A_des, b)[0]
        #x = scipy.linalg.solve(A_des.T@A_des,A_des.T@(b))
        x= np.linalg.solve(A_des.T@A_des,A_des.T@b)
        

    return x

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