如何在Python中高效执行10000次lstsq计算？

我正在编写代码，将总计10000个粒子投影至MLS（移动最小二乘）曲面上。该任务需要调用最小二乘函数10000次，每次计算所用矩阵的形状为(N,6)（N根据粒子位置不同，取值在200到1000之间）。

当前我使用torch.linalg.lstsq进行计算，每次调用耗时约0.5秒，因此寻求更高效的实现方案，若能提供代码示例或推荐库将极为有用。

已尝试的方法 #

• 使用numpy.linalg.lstsq搭配for循环，耗时约1.3秒。
• 使用scipy.linalg.lstsq并指定gelsy LAPACK驱动，搭配for循环，耗时约1.3秒。
• SVD方法，耗时约1.5秒，代码如下：

u, s, v = np.linalg.svd(A, full_matrices=False)
uTb = np.einsum('ijk,ij->ik', u, b)
c = np.einsum('ijk,ij->ik', v, uTb / s)
return c

• 使用np.linalg.solve求解A.T@Ax=A.T@b，耗时约1.5秒，代码如下：

ATA = np.einsum('ijk,ijl->ikl', A, A)
ATb = np.einsum('ijk,ij->ik', A, b)
c = np.linalg.solve(ATA, ATb)
return c

• 多线程方法：因代码其他部分使用Taichi库，尝试多线程时会初始化多个Taichi后端，导致电脑死机，代码如下：

from multiprocessing import Pool
with Pool() as pool:
      c = pool.starmap(calc_lstsq, [(A[i], b[i]) for i in range(b.shape[0])])
return np.asarray(c)

内容的提问来源于Stack Exchange，提问作者konbu