在大型阵列中使用广播、转置和重塑的 Numpy 速度效率答案

【问题标题】：Numpy speed efficiency using broadcasting, transpose and reshape in large size array在大型阵列中使用广播、转置和重塑的 Numpy 速度效率
【发布时间】：2021-12-09 04:46:36
【问题描述】：

有没有办法加快下面这行代码：

desired_channel=32
len_indices=50000
fast_idx = np.broadcast_to(np.arange(desired_channel)[:, None], (desired_channel, len_indices)).T.reshape(-1)

谢谢。

【问题讨论】：

那你有什么建议？
Numba 应该有助于加速使用普通循环的代码，尽管生成的代码将非常低级。
@user2950931 如果有任何答案解决了您的问题，请考虑accepting it，如果有帮助，请单击复选标记或点赞。这向更广泛的社区表明您已经找到了解决方案，并为回答者和您自己提供了一些声誉。没有义务这样做。

标签： numpy performance reshape array-broadcasting

【解决方案1】：

最后一行代码等于np.tile(np.arange(desired_channel), len_indices)。

在我的机器上，np.tile 的性能与许多 Numpy 调用一样受到操作系统（页面错误）、内存分配器和内存吞吐量的限制。有两种方法可以克服这个限制：不分配/填充临时缓冲区，使用较短的类型在内存中生成较小的数组，例如@987654323 @ 或 np.uint16 关于您的需求。

由于np.tile 函数没有out 参数，Numba 可用于生成快速替代函数。这是一个例子：

import numba as nb

@nb.njit('int32[::1](int32, int32, int32[::1])', parallel=True)
def generate(desired_channel, len_indices, out):
    for i in nb.prange(len_indices):
        for j in range(desired_channel):
            out[i*desired_channel+j] = j
    return out

desired_channel=32
len_indices=50000
buffer = np.full(desired_channel * len_indices, 0, dtype=np.int32)
%timeit -n 200 generate(desired_channel, len_indices, fast_idx)

以下是性能结果：

Original code: 1.25 ms
np.tile:       1.24 ms
Numba:         0.20 ms

【讨论】：

【解决方案2】：

我是 jax 库的新手。我在 Colab TPU 上使用以下代码通过 jax one 比较了您的代码：

import numpy as np
from jax import jit
import jax.numpy as jnp
import timeit

desired_channel=32
len_indices=50000

def ex_():
    return np.broadcast_to(np.arange(desired_channel)[:, None], (desired_channel, len_indices)).T.reshape(-1)
%timeit -n1000 -r10 ex_()

@jit
def exj_():
    return jnp.broadcast_to(jnp.arange(desired_channel)[:, None], (desired_channel, len_indices)).T.reshape(-1)
%timeit -n1000 -r10 exj_()

在我的一项努力中，结果如下：

1000 个循环，10 个循环中的最佳：每个循环 901 µs
1000 个循环，最好的 10 个：每个循环 317 µs

通过这种方式，jax 可以将您的代码加速大约两到三倍。

【讨论】：