沿轴应用函数的最快方法答案

【问题标题】：Fastest way to apply function along axes沿轴应用函数的最快方法
【发布时间】：2023-01-20 03:24:48
【问题描述】：

在时间关键的代码片段中，我需要沿张量的不同轴应用一个函数并对结果求和。一个奇特的特征是张量 (ns_test) 的轴数可以很大。我想出了两个实现，我将当前轴 (moveaxis) 移动到第零个 (h_zero) 或最后一个 (h_last) 位置，应用函数，然后将轴移回。我不确定这是最好的方法。

import numpy as np
import time

def h_last(state, km, ns):
    new_state = np.zeros_like(state)
    for i in range(ns):
        a = np.moveaxis(state, i+1, -1).copy()
        for k in range(km):
            a[..., k] = (k+0.5) * a[..., k]
        new_state += np.moveaxis(a, -1, i+1)
    return new_state

def h_zero(state, km, ns):
    new_state = np.zeros_like(state)
    for i in range(ns):
        a = np.moveaxis(state, i+1, 0).copy()
        for k in range(km):
            a[k, ...] = (k+0.5) * a[k, ...]
        new_state += np.moveaxis(a, 0, i+1)
    return new_state

# ==================== init ============================
km_test  = 4        
ns_test = 7
nreps = 100
dims = tuple([ns_test] + [km_test] * ns_test)    
y= np.random.rand(*dims)
    
# =================== first run =============================
tic = time.perf_counter()
for i in range(nreps):
    yy = h_last(y, km_test, ns_test)
toc = time.perf_counter()
print(f"Run time h_last {toc - tic:0.4f} seconds")

# =================== second run =============================
tic = time.perf_counter()
for i in range(nreps):
    yyy = h_zero(y, km_test, ns_test)
toc = time.perf_counter()
print(f"Run time h_zero {toc - tic:0.4f} seconds")

print(np.linalg.norm(yy-yy))

我有点惊讶第零轴的性能更好（我认为 python 在内部使用 C 顺序进行存储）。但我的主要问题是如何进一步加速代码？我查看了apply_along_axis，但这似乎很慢。

【问题讨论】：

为什么要复制 moveaxis 结果？ moveaxis 的全部意义不就是要快，因为它只是一个视图吗？
@chrslg 因为我在a[..., k] = (k+0.5) * a[..., k]这一行修改了a并且因为我不想通过这个操作破坏state。
是的，我明白了。但这使整个事情变得非常缓慢。因为你为每个轴复制。也许你应该复制state并尽可能长时间地工作。
@chrslg 我考虑过，但无法提出有效的解决方案。
因此，简而言之，从“C 顺序/缓存”的角度来看，a[k,...] *= 操作比a[...,k] *= 操作快。我想这就是你所看到的。如果我总结一下我对这个性能问题的所有抱怨：如果你没有处理 move_axis 结果的副本，那么，你使用的假轴位置是什么根本无关紧要。但既然你是，zero 比 last 更好是很正常的，因为重要的是循环的最内层索引是最后一个轴（按 C 顺序）。这里循环的最内层索引出现在numpy的隐式循环中a[k, ...] *= ...

标签： python numpy optimization tensor

【解决方案1】：

正如 cmets 中所讨论的那样，moveaxis 很快。这只是一个观点。因此，当您有如此多的轴需要一个循环来迭代轴号并在给定轴上执行操作时，这是一种有趣的工作方法。

减慢计算速度的是您在整个过程中对数组进行的多次复制。另外，正如在 cmets 中所说，由于“内存缓存和 C 顺序”的考虑，也是这个副本使得最后一个轴的版本稍微慢一些。我不会在这里开发我在 cmets 中所说的内容，因为它非常附属（我们说的是 20% 的性能损失。不可忽略，但与您的实际问题相比确实没有什么）

因此，更快地执行您描述的操作的一种方法是避免这些副本

例如，这是一段代码，大致使用与您相同的技巧（使用moveaxis）

tic = time.perf_counter()
kh=np.arange(0.5, km_test).reshape([1]*ns_test+[-1])

Y=np.zeros_like(y)
for i in range(ns_test):
    Y += y*np.moveaxis(kh, -1, i+1)
toc = time.perf_counter()
print(f"Run time arange {toc - tic:0.4f} seconds")

print(np.linalg.norm(yy-Y)) # Note that in your code you compare yy with yy. You probably meant yyy for one of the two yy.

结果在我的电脑上

Run time h_last 1.2403 seconds
Run time h_zero 0.9995 seconds
Run time arange 0.0061 seconds
0.0

所以结果是一样的。但计算速度要快 150-200 倍。

然而，我和你一样经常使用moveaxis。当然，在较小的阵列上，因为我选择仅在 arange 阵列上使用它。移动moveaxis 成本与轴的数量成正比，而不是数据的大小（这只是一些步幅和类似的东西。没有数据被移动。这就是重点）。此外，由于这个arange，我避免了一个 for 循环。但这可能不是解释大部分性能比的原因。因为那是一个外部 for 循环（在性能方面重要的是内部循环，在 numpy * 操作中）。所以，剩下来解释性能比的是副本，我会说。

请注意，例如，如果我这样做

Y += np.moveaxis(np.moveaxis(y, i+1, -1)*kh, -1, i+1)

（即移动 y 轴而不是移动 kh 轴。这迫使我在使用结果之前将轴向后移动，就像你所做的那样）而不是我之前的 Y+=... 行，它也是相同的结果，并且相同一种性能，同时更接近您的方法。 moveaxis 大致不花钱。所以我选择移动轴 kh 更符合我的口味，但这并不是导致性能差异的原因。

【讨论】：