在Python中加速一个与变量相关的矩阵求逆计算。

Question

我有一个迭代问题，在每次迭代中，我需要执行一系列矩阵操作，我想加快速度。大部分计算涉及常量变量，除了一个可变的变量。 我的问题是：是否可能构建这个操作的“骨架”，以加快第一次迭代后的计算速度？（类似于CVXPY优化问题的参数热启动）。 我需要处理的代码目前在类中的专用函数中，类似于这样：

# some constants definitions (in a class)
m = 5501
Y = np.random.rand(10, m)
U = np.random.rand(75, m)
sig_on = 0.15
sig_off = 0.25
a = 15
L = 25

def obtain_A_B(x):
   # for readability I omit internal variable declarations as a = self.a
   varying_part = (a * sig_on / np.linalg.norm(x)) ** 2 + L * sig_off
   F = varying_part * np.identity(m) + Y.T @ Y
   Fi = np.linalg.inv(F)
   UFiUTi = np.linalg.inv(U @ Fi @ U.T)
   A = (Fi - Fi @ U.T @ UFiUTi  @ U @ Fi) @ Y.T
   B = Fi @ U.T @ UFiUTi 
   return A, B

x = np.random.rand(10)
result = obtain_A_B(x)

我对加快 A 和 B 的计算很感兴趣，'varying_part' 返回一个标量，它会随着给定 x 的每次迭代而变化。 由于真正耗时的部分是计算逆矩阵 Fi，只要加速这个计算就可以帮助提高效率。

我已经将逆矩阵移动到专用变量中（如上所述），以减少需要计算的次数 - 这将执行时间从大约41秒降至大约10秒 - 但最好能够让执行速度更快（这就是“操作骨架”想法的原因）。

- Andrea

---

首先，您计算了 np.linalg.inv(U @ Fi @ U.T) 两次 - 实际上，U.T @ np.linalg.inv(U @ Fi @ U.T) 可能可以被缓存。 - AKX

但通常情况下，在尝试优化之前，先使用分析器找出哪里是慢的部分。 - AKX

耗时最长的部分仍然是求逆 - 6秒，其余部分约为4秒，并且计算np.linalg.inv(U @ Fi @ U.T)需要约0.093秒（如预期所示，不是问题所在）。我已经在Matlab中解决了这个问题，并彻底调查了时间：随着问题规模的增加（增加m），主要问题是求逆。 - Andrea

你可以尝试使用numba，因为该函数只涉及Numpy操作... - AKX

我修改了代码，使其具有可重现性，同时保持相同的逻辑（如果有任何错误，请随意更改）。至于Numba，在这里并不是非常有用，除非编写比优化的BLAS实现更快的代码，这是可能的，但远非易事。 - Jérôme Richard

Answer 1

基本优化

首先，A 中使用了 B 表达式，因此可以先计算 B 并在 A 表达式中重复使用。

此外，varying_part * np.identity(m) + Y.T @ Y 可以更高效地计算。事实上，varying_part * np.identity(m) 创建了两个大的临时数组，填充速度较慢。这是因为如今的 RAM 与 CPU 性能相比较慢（参见"Memory Wall"）。你只需要更新该数组的对角线即可。

另外，计算 B @ (U @ Fi) 要比 (B @ U) @ Fi（等同于 B @ U @ Fi）更快，尽管它们在解析上是等价的，因为矩阵乘法是结合的（虽然数值上略有不同）。

以下是优化后的实现：

def obtain_A_B_opt(x):
   varying_part = (a * sig_on / np.linalg.norm(x)) ** 2 + L * sig_off
   F = Y.T @ Y
   np.fill_diagonal(F, F.diagonal()+varying_part)
   Fi = np.linalg.inv(F)
   tmp = Fi @ U.T
   UFiUTi = np.linalg.inv(U @ tmp)
   B = tmp @ UFiUTi 
   A = (Fi - B @ (U @ Fi)) @ Y.T
   return A, B

这个优化使我的机器上的代码快了55%。近90%的时间花费在np.linalg.inv(F)函数调用中。

坏消息是矩阵求逆在CPU上很难进行优化。Numpy使用的默认OpenBLAS库执行此计算是次优的，但它相对较好，因为它使用相对优化的本机C代码并行运行。此外，改变对角线会强烈影响计算，因此我不认为可以编写逆的增量计算。实际上，将F[0,0]增加1会导致高斯约旦算法中所有行的权重都必须重新计算（因此所有后续权重计算）。

如果不减少浮点精度或算法的准确性（例如近似），很难编写更快的代码。如果您拥有服务器GPU或者可以减少精度，GPU可以显著提高速度。

---

基于GPU的实现

可以使用GPU使该代码在目标平台和输入浮点类型方面更快。离散GPU tend tend to have a significantly bigger computational power than mainstream CPUs. 然而，在主流PC GPU上双精度不成立。事实上，PC GPU针对简单精度浮点数据进行了优化（常用于游戏和3D应用程序）。对于快速双精度计算，需要服务器端GPU。此外，应该记住，CPU和GPU之间的数据传输是昂贵的，但矩阵求逆是瓶颈，并且它在O(N**3)中运行，因此数据传输在这里可以忽略不计。Cupy可用于轻松编写GPU实现（适用于Nvidia GPU）。以下是生成的GPU代码：

import cupy as cp

def obtain_A_B_gpu_opt(x):
   varying_part = (a * sig_on / np.linalg.norm(x)) ** 2 + L * sig_off

   # Send data to the GPU
   device_Y = cp.array(Y)
   device_U = cp.array(U)

   # GPU-based computation
   F = device_Y.T @ device_Y
   cp.fill_diagonal(F, F.diagonal()+varying_part)
   Fi = cp.linalg.inv(F)
   tmp = Fi @ device_U.T
   UFiUTi = cp.linalg.inv(device_U @ tmp)
   B = tmp @ UFiUTi 
   A = (Fi - B @ (device_U @ Fi)) @ device_Y.T

   # Get data back from the GPU
   return A.get(), B.get()

请注意，Cupy在首次执行时可能会非常缓慢，但之后应该会明显加快（这对于迭代循环来说是可以接受的）。

---

实验结果

下面是在我的机器上使用i5-9600KF CPU和Nvidia 1660 Super GPU（主流离散式个人电脑GPU）得到的结果：

double-precision:
 - obtain_A_B:          4.10 s
 - obtain_A_B_opt:      2.64 s
 - obtain_A_B_gpu_opt:  3.03 s

simple-precision:
 - obtain_A_B:          3.17 s
 - obtain_A_B_opt:      2.52 s
 - obtain_A_B_gpu_opt:  0.28 s  <----

可以看出，基于GPU的实现在使用单精度浮点输入时速度显着更快。这是因为我的GPU被优化过。在计算服务器上，obtain_A_B_gpu_opt 应该比其他实现要快得多。请注意，在我的GPU上，cp.linalg.inv 仍然占用了90%以上的时间。

- Jérôme Richard