重构算法以避免大型 Numpy 数组的低效切片

Question

我有一个工作算法来分析实验数据集。该算法由两个主要函数组成。第一个将一个大数组作为其输入之一，returns 一个通常不适合内存的中间 3D 复值数组。为此，我使用 Numpy 的 memmap 将这个数组保存在磁盘上。对于更大的数据集，第二个函数开始花费更多时间并且它似乎与内存访问有关。在某些情况下，如果计算持续 20 分钟，则将 n2 增加 50% 会导致计算需要将近 24 小时。

剥离整个程序的大约 99% 后，一个最小的工作示例如下所示：

import numpy as np

n1, n2 = 10, 100

Nb = 12
Nf = int(n2 / Nb)

X = np.random.rand(n1, n2)

#   Function 1: Compute X2 from X1 and save on disk
X2 = np.memmap("X2.dat", shape=(Nb, n1, Nf), dtype=np.complex128, mode='w+')

for n in range(Nb):
    Xn = X[:, n*Nf:(n+1)*Nf]
    X2[n,:,:] = np.fft.fft(Xn, axis=1)

X2.flush()
del X2

#   Function 2: Do something with slices of X2
X2 = np.memmap("X2.dat", shape=(Nb, n1, Nf), dtype=np.complex128, mode='r')

for k in range(Nf):
    Y = np.pi * X2[:,:,k]  #   <- Problematic step
    # do things with Y...

n1、n2、Nb 和 Nf 的值通常更大。

如您所见，在函数 1 中，X2 使用其第一个索引进行填充，根据的说法，就速度而言，这是最有效的方法。我的问题出现在函数 2 中，我需要通过在其第三维上对其进行切片来处理 X2。这是算法的本质要求。

我想找到一种方法来重构这些函数以减少执行时间。我已经改进了算法的某些部分。例如，因为我知道 X 只包含实数值，所以如果我忽略它的共轭值，我可以将 X2 的大小减少 2。但是， X2 的切片仍然存在问题。（实际上，我只是在写这篇文章时了解到 np.fft.rfft，这对我绝对有帮助）。

您是否看到我可以重构函数 2（and/or 函数 1）以便更有效地访问 X2 的方法？

更新

我在我最大的一个数据集上测试了，结果表明第一个优化，将 Nf 移动到轴 1，总体上比将它移动到轴 0 稍微快一些。

下面的三个图表分别显示了总执行时间、函数 1 花费的时间（不包括 X2.flush()）和函数 2 花费的时间的分析结果。在 x 轴上，Nr 是与 n2 和 Nb 成比例的值。我使用 Numpy 的 rfft() 对我的初始代码进行了优化和修改后的代码测试，同时也进行了两种优化。

使用我的初始代码，选择。 1 是更好的选择，Nr=12 的总时间减少了一个数量级以上。使用 rfft() 几乎可以减少另一个数量级的时间，但在这种情况下，两种优化是等效的（一切都适合可用的 RAM，因此交换数组轴的时间减少是最小的）。但是，这将使更有效地处理更大的数据集成为可能！

Answer 1

一个简单的优化是交换最后两个轴，这不应该改变函数 1 的速度（假设来自转置的 out-of-order 内存访问与磁盘访问相比可以忽略不计）但应该使函数2 运行更快，原因与您链接的问题中讨论的相同：

#   Function 1: Compute X2 from X1 and save on disk
X2 = np.memmap("X2.dat", shape=(Nb, Nf, n1), dtype=np.complex128, mode='w+')

for n in range(Nb):
    Xn = X[:, n*Nf:(n+1)*Nf]
    X2[n,:,:] = np.fft.fft(Xn, axis=1).T  # swap axes so Nf changes slower

X2.flush()
del X2

#   Function 2: Do something with slices of X2
X2 = np.memmap("X2.dat", shape=(Nb, Nf, n1), dtype=np.complex128, mode='r')

for k in range(Nf):
    Y = np.pi * X2[:,k,:]
    # do things with Y...

您可以通过将 Nf 移动到轴 0:

以降低功能 1 的速度为代价使功能 2 更快

#   Function 1: Compute X2 from X1 and save on disk
X2 = np.memmap("X2.dat", shape=(Nf, Nb, n1), dtype=np.complex128, mode='w+')

for n in range(Nb):
    Xn = X[:, n*Nf:(n+1)*Nf]
    X2[:,n,:] = np.fft.fft(Xn, axis=1).T  # swap axes so Nf changes slower

X2.flush()
del X2

#   Function 2: Do something with slices of X2
X2 = np.memmap("X2.dat", shape=(Nf, Nb, n1), dtype=np.complex128, mode='r')

for k in range(Nf):
    Y = np.pi * X2[k,:,:]
    # do things with Y...

如果 X2 的读取次数多于写入次数，则使用此版本可能有意义。此外，函数 1 的减速应该随着 n1 变大而变小，因为连续的块更大。

将数据文件存储在硬盘上 n1, n2, Nb = 1000, 10000, 120，我的时间是

function 1, original:          1.53 s ± 41.9 ms
function 1, 1st optimization:  1.53 s ± 27.8 ms
function 1, 2nd optimization:  1.57 s ± 34.9 ms

function 2, original:          111 ms ± 1.2 ms
function 2, 1st optimization:  45.5 ms ± 197 µs
function 2, 2nd optimization:  27.8 ms ± 29.7 µs

重构算法以避免大型 Numpy 数组的低效切片

Refactoring an algorithm to avoid inefficient slicing of a large Numpy array

python

memory

optimization

performance

numpy