numpy 切片在函数中的性能不佳

Question

我有这样的示例代码

import numpy as np

nbins = 1301
init = np.ones(nbins*2+1)*(-1)
init[0] = 0
init[1::2] = 0
z0 = np.array(init, dtype=np.complex128, ndmin=1)
initn = z0.view(np.float64)
deltasf = np.linspace(-10, 10, nbins)
gsf = np.linspace(1, 2, nbins)

def jacobian_mbes_test(Y, t):
     leny = len(Y)
     J = np.zeros((leny, leny))
     J[0,0] = -10
     J[0,1] = 0
     J[0, 2::4] = gsf
     J[1,0] = 0
     J[1,1] = -10
     J[1,3::4] = gsf
     J[2::4, 0] = gsf*Y[4::4]
     J[3::4, 0] = gsf*Y[5::4]
     J[4::4, 0] = -4*gsf*Y[2::4]
     J[2::4, 1] = -gsf*Y[5::4]
     J[3::4, 1] = gsf*Y[4::4]
     J[4::4, 1] = -4*gsf*Y[3::4]
     J[(range(2, leny, 4), range(2, leny, 4))] = -0.8
     J[(range(3, leny, 4), range(3, leny, 4))] = -0.8
     J[(range(4, leny, 4), range(4, leny, 4))] = -0.001
     J[(range(5, leny, 4), range(5, leny, 4))] = -0.001
     J[(range(2, leny, 4), range(3, leny, 4))] = deltas
     J[(range(3, leny, 4), range(2, leny, 4))] = -deltas
     J[(range(2, leny, 4), range(4, leny, 4))] = gsf*Y[0]
     J[(range(2, leny, 4), range(5, leny, 4))] = -gsf*Y[1]
     J[(range(3, leny, 4), range(4, leny, 4))] = gsf*Y[1]
     J[(range(3, leny, 4), range(5, leny, 4))] = gsf*Y[0]
     J[(range(4, leny, 4), range(2, leny, 4))] = -4*gsf*Y[0]
     J[(range(4, leny, 4), range(3, leny, 4))] = -4*gsf*Y[1]
     return J

使用scipy.integrate.odeint

计算一组微分方程的雅可比行列式

不幸的是函数jacobian_mbes_test的性能不是很好。 %timeit jacobian_mbes_test(initn, 0) 给出了 12.2 毫秒，即使我试图用切片（？）有效地完成所有事情。

将函数更改为（因为我猜更复杂的索引分配需要最多时间）

def jacobian_mbes_test2(Y, t):
     leny = len(Y)
     J = np.zeros((leny, leny))
     J[0,0] = -10
     J[0,1] = 0
     J[0, 2::4] = gsf
     J[1,0] = 0
     J[1,1] = -10
     J[1,3::4] = gsf
     J[2::4, 0] = gsf*Y[4::4]
     J[3::4, 0] = gsf*Y[5::4]
     J[4::4, 0] = -4*gsf*Y[2::4]
     J[2::4, 1] = -gsf*Y[5::4]
     J[3::4, 1] = gsf*Y[4::4]
     J[4::4, 1] = -4*gsf*Y[3::4]
     return J

将这个时间减少到 4.6 毫秒，这仍然不是很好。有趣的是，如果我在

这样的函数之外对这些赋值进行计时

J = np.zeros((len(initn), len(initn))
%timeit J[4::4, 1] = -4*gsf*initn[3::4]

我最多得到 10µs，所以我希望第二个函数最多（！）像 100µs 一样。是否有一些我不知道的函数在内存中进行复制？

是否有更有效的方法来为某些索引赋值？

Answer 1

问题来自矩阵大小 (206 MiB)。实际上，矩阵非常大并且（实际上）每次调用函数时都用零填充。假设所有值都将在 12.2 毫秒内物理写入内存，吞吐量将为 5206**2*8/12.2e-3/1024**3 = 16.5 GiB/s，这对于顺序工作负载来说非常好（尤其是在通常的个人计算机上勉强达到连续 25 GiB/s）。实际上，内存并未设置为零，大部分计算时间来自 虚拟内存.

的管理

np.zeros((leny, leny)) 内部请求您的操作系统 (OS) 保留一些 space 并将其初始化为零。您的 OS 实际上会分配许多页（通常为 4 KiB 的块）并将它们标记为稍后设置为零（延迟初始化）。因此，这个过程非常便宜。但是，稍后执行J[0,0] = -10时，它会在给定的first-touched page中写入需要初始化的数据。这称为 页面错误 。这个过程是昂贵的（尤其是顺序的），因为它停止你的过程的执行，在物理内存中找到一些地方，填充整个页面，然后恢复你的过程，每个页！在您的情况下，许多要初始化的数组单元格会导致页面被填充，这比实际填充单元格要昂贵得多。

CPU 缓存 在此代码中也很重要，因为当在 RAM 中完成计算时，页面填充会慢得多。当数据量太大而无法放入最后一级缓存（通常为 1-64 MiB）时，就会发生这种情况。

内存访问模式 强烈影响相当大的数据结构的性能（即不适合 CPU 缓存）。事实上，随机访问模式或大跨度的跨步访问对性能来说确实很糟糕，因为处理器无法轻易预测要提前加载的 RAM 块的位置（更不用说 RAM 延迟非常大）。连续访问要快得多。

在比较简单的情况下测量时间时，应考虑所有这些因素。

查看虚拟内存影响的一种简单方法是使用 np.full((leny, leny), 0.0) 而不是 np.zeros(leny, leny)。这在我的机器上慢了 3 倍。在这种情况下，所有页面都需要填写。您还可以创建测量时间来多次填充 J 矩阵。第一次应该慢得多（在我的机器上慢 2.2 倍）。

一个解决方案是分配和初始化数组一次，fill/recycle 只分配和初始化一些值，但这有点棘手。在最坏的情况下，大部分矩阵都需要归零，这在您的情况下太慢了，因为它会受到内存吞吐量的限制。

一个简单的解决方案是使用space矩阵。 Space 矩阵可以使用模块 scipy.sparse 创建。它们的计算速度通常较慢，但如果有很多零（这是您的情况），它们的计算速度会更快且更紧凑。这是一个例子：

from scipy.sparse.lil import lil_matrix

def jacobian_mbes_test(Y, t):
     leny = len(Y)
     J = lil_matrix((leny, leny))
     J[0,0] = -10
     J[0,1] = 0
     J[0, 2::4] = gsf
     J[1,0] = 0
     J[1,1] = -10
     J[1,3::4] = gsf
     J[2::4, 0] = gsf*Y[4::4]
     J[3::4, 0] = gsf*Y[5::4]
     J[4::4, 0] = -4*gsf*Y[2::4]
     J[2::4, 1] = -gsf*Y[5::4]
     J[3::4, 1] = gsf*Y[4::4]
     J[4::4, 1] = -4*gsf*Y[3::4]
     J[(range(2, leny, 4), range(2, leny, 4))] = -0.8
     J[(range(3, leny, 4), range(3, leny, 4))] = -0.8
     J[(range(4, leny, 4), range(4, leny, 4))] = -0.001
     J[(range(5, leny, 4), range(5, leny, 4))] = -0.001
     J[(range(2, leny, 4), range(3, leny, 4))] = deltas
     J[(range(3, leny, 4), range(2, leny, 4))] = -deltas
     J[(range(2, leny, 4), range(4, leny, 4))] = gsf*Y[0]
     J[(range(2, leny, 4), range(5, leny, 4))] = -gsf*Y[1]
     J[(range(3, leny, 4), range(4, leny, 4))] = gsf*Y[1]
     J[(range(3, leny, 4), range(5, leny, 4))] = gsf*Y[0]
     J[(range(4, leny, 4), range(2, leny, 4))] = -4*gsf*Y[0]
     J[(range(4, leny, 4), range(3, leny, 4))] = -4*gsf*Y[1]
     return J

这在我的机器上快了 2 倍。有许多不同类型的稀疏矩阵，每一种都适用于一些特定的用例。可能有比 LIL 更好的稀疏矩阵类型。

LIL spaces 矩阵对于具有动态结构的 space 矩阵非常好。 CSR 非常适合静态的。因此，例如，您可以提前构建一个 CSR 矩阵，在您计划写入一些数据的位置填充 NaN 值，然后复制它并在需要时用新值填充新的 space 矩阵。这在我的机器上快了 10-15%。使用更好的访问模式可能会进一步减少时间。

numpy 切片在函数中的性能不佳

Bad performance of numpy slicing in function

python

performance

numpy

slice

numpy-slicing