NumPy 为每个子张量沿轴指定函数应用？答案

【问题标题】：NumPy apply along axis specifying function for each subtensor?NumPy 为每个子张量沿轴指定函数应用？
【发布时间】：2021-08-07 11:14:29
【问题描述】：

假设有一个形状为(N, M) 的矩阵A 和一些dtype。然后，有一个numpy 函数B 的向量，形状为(N, )。因为这可能很重要，对于我的具体情况，N 的范围以千为单位（例如 4000），M 的范围以单位为单位（例如 9）。

示例（N=5,M=3）：

A =
[
    [0.1, 0.2, 0.3],
    [4.02, 123.4, 534.65],
    [2.32, 22.0, 754.01],
    [5.41, 23.1, 1245.5],
    [6.07, 0.65, 22.12],
]

B = [
    np.min,
    np.max,
    np.argmin,
    np.argmin,
    np.min
]

B[i] 将始终为 np.min、np.max、np.argmin 或 np.argmax。

我想对axis=1 执行操作，特别是对A[i, :] 执行操作B[i]，产生以下输出：

[ 0.1, 534.65, 0, 0, 0.65 ]

这样的事情在 NumPy 中可以实现吗？

一种可能的解决方案是根据B 上的等价组拆分A，沿axis=1 应用单个操作，重新索引组并将它们合并回来，但这听起来效率很低。如果可能的话，我想避免这种情况。

如果我的提议对你来说听起来很荒谬，请告诉我。

【问题讨论】：

如果您关心效率，那么与 A 中的列数相比，函数的数量是否小很重要。是吗？
您提出的结果沿axis=0（逐行）应用
要么是这样，要么是对zip(A,B) 的简单迭代。 B 是对象 dtype，本质上是一个列表。
@user2640045 抱歉，我应该在N >> M 中添加该信息。 N 以千为单位，M 以单位为单位。
@SamarthBhatia 我作为示例向您展示的结果实际上与axis=1 一起应用。第一个轴是行，第二个轴是列。

标签： python numpy

【解决方案1】：

为什么这对您来说不可行？：

A = np.array(A) # A originally is what you have given
B = [...] # what you have given

for i in range(len(B)):
    print( np.apply_along_axis(B[i], 0, A[i,:]) )

注意：它比递归执行B[i](A[i,:]) 更快。检查此link（这是numpy 文档）

【讨论】：

apply_along_axis 添加了什么？ B[i](A[i]) 更简单。
@hpaulj 我已经编辑了我的答案以包含来自 numpy 文档的参考
这样会不会非常低效？您正在使用 for 循环分别调用所有这些函数。
@CaptainTrojan np.apply_over_axes 存在，它将 a 函数应用于所有轴。但是，我不知道如何使用函数列表来做到这一点。
@SamarthBhatia 另一个注意事项，aai(e) 用于内置 numpy 函数没有意义。您可以只获取该函数并调用它，而不是将其提供给另一个函数。因此，您可以简单地执行B[i](A[i,:]) 并节省时间。感谢您尝试回答我的问题！

【解决方案2】：

您的 2 个列表：

In [162]: A =[
     ...:     [0.1, 0.2, 0.3],
     ...:     [4.02, 123.4, 534.65],
     ...:     [2.32, 22.0, 754.01],
     ...:     [5.41, 23.1, 1245.5],
     ...:     [6.07, 0.65, 22.12],
     ...: ]
     ...: 
     ...: B = [
     ...:     np.min,
     ...:     np.max,
     ...:     np.argmin,
     ...:     np.argmin,
     ...:     np.min
     ...: ]

简单的列表理解：

In [164]: [b(a) for b,a in zip(B,A)]
Out[164]: [0.1, 534.65, 0, 0, 0.65]

请注意，列表元素是浮点数和整数、值和索引的混合体。

In [165]: timeit [b(a) for b,a in zip(B,A)]
61.1 µs ± 918 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

事先将A 设为数组会有所帮助。使B 成为一个数组（object dtype）并没有帮助：

In [167]: %%timeit A1=np.array(A)
     ...: [b(a) for b,a in zip(B,A1)]
35.6 µs ± 36.5 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

关于您在B 中收集类似值函数的想法，这是一个开始：

跳过收集B的类似元素的逻辑，这些是组：

In [169]: idx = [[0,4],[1],[2,3]]
In [171]: arr = np.array(A)       # has to be array for multirow indexing
In [172]: [B[i[0]](arr[i], axis=1) for i in idx]
Out[172]: [array([0.1 , 0.65]), array([534.65]), array([0, 0])]
In [173]: timeit [B[i[0]](arr[i], axis=1) for i in idx]
49.2 µs ± 1.79 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

这里不是更快，但对于具有许多重复的大型数组可能会节省一些时间。

要重建有序列表，请将其更改为分配给对象 dtype 数组的循环：

In [181]: res = np.empty(len(B), object)
     ...: for i in idx:
     ...:     res[i] = B[i[0]](arr[i], axis=1)
     ...: 
In [182]: res
Out[182]: array([0.1, 534.65, 0, 0, 0.65], dtype=object)
In [183]: %%timeit
     ...: res = np.empty(len(B), object)
     ...: for i in idx:
     ...:     res[i] = B[i[0]](arr[i], axis=1)
     ...: 
71.5 µs ± 48 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

收集B 之类的元素的一种方法是使用dict：

In [185]: dd = defaultdict(list)
In [186]: for i,b in enumerate(B): dd[b].append(i)
...
In [189]: list(dd.values())
Out[189]: [[0, 4], [1], [2, 3]]
In [190]: res = np.empty(len(B), object)
     ...: for k in dd:
     ...:     i = dd[k]
     ...:     res[i] = k(arr[i], axis=1)
     ...: 
In [191]: res
Out[191]: array([0.1, 534.65, 0, 0, 0.65], dtype=object)

【讨论】：

感谢您的回答！对不起，如果这不感恩，但基准在这里有点不合适。确实，我准确地发布了A 和B 的这些示例，但我还提到N 非常大，比M 大得多，而且肯定比5 大得多。我认为性能差异会更明显，因为我不能以这些不确定的结果为基础。 +1 的努力，但我不能接受这一点，因为你唯一的贡献是已经提出的解决方案的基准，不幸的是 N 的值太小。
我将真实世界的基准测试留给您。变量太多 - A 的两个维度，B 的副本等。分组，表面上是更多的工作，但可以节省时间。也许是 2 倍，但我怀疑更多。但是为什么首先要混合这样的动作呢？只需将 min 应用于一个值数组，将 aargmax 应用于另一个值等等。B 列表有些人为。
如果您有兴趣，可以查看我自己的答案。这个因素实际上倾向于略高于8 并保持在7 和8 之间。当然，用例难以想象是对的，但这样做的真正动机是 a) 泛化和 b) 强化学习技巧。例如，我有一个完整的计算数组 Q-values 但需要它们的 max 或 min w.r.t。 axis=1 根据当前玩家（由B 或可1:1 映射到B 的东西给出）是白色还是黑色。
我实际上所做的是将min/max 和argmin/argmax 分开，因为我偷偷知道B 永远不会包含foo 和argfoo 在在同一时间，只要B[i] 是(arg)min 时，只需将A[i] 乘以-1，然后在整个过程中调用(arg)max，并将标志改回来。当我意识到我可以做到这一点时，这个问题变得有点哲学：D 希望你玩得开心，呵呵。

【解决方案3】：

为了补充@hpaulj 所做的基准测试，我创建了这个脚本：

import numpy as np
import matplotlib.pyplot as plt
from time import perf_counter_ns

COUNT = 100
M = 9

x = []
y_1 = []
y_2 = []

fns = [np.min, np.max, np.argmin, np.argmax]

for n in range(2, 200, 1):
    print(n)
    A = np.random.rand(n, M)
    B = np.random.choice(fns, n)

    x.append(n)

    y1acc = []
    y2acc = []

    for i in range(COUNT):
        start = perf_counter_ns()
        res1 = [b(a) for b, a in zip(B, A)]
        y1acc.append(perf_counter_ns() - start)

        start = perf_counter_ns()
        groups = {fn: [] for fn in fns}

        for j, fn in enumerate(B):
            groups[fn].append(j)

        res2 = np.empty(n, object)

        for fn, indices in groups.items():
            res2[indices] = fn(A[indices], axis=1)

        y2acc.append(perf_counter_ns() - start)

        assert np.array_equal(res1, res2)

    y_1.append(sum(y1acc) / COUNT)
    y_2.append(sum(y2acc) / COUNT)

plt.plot(x, y_1, label="zip")
plt.plot(x, y_2, label="groups")
plt.legend()
plt.show()

并观察到这些结果：

显然，拆分和合并小数组的开销远大于由 O(1) 调用 np.foo 而不是 O(N) 导致的 NumPy 加速。

对于较大的 N (for n in range(100, 5000, 100))：

我们清楚地观察到执行这些操作所期望的线性行为，然而，分裂合并方法中的常数要小得多。请注意，此基准测试不是在受控条件下执行的，但引入的噪声可以忽略不计。

【讨论】：