python - 在numpy中将同一行堆叠到堆叠的二维矩阵的最佳方法是什么？

Question

我正在寻找一种方法来有效地将同一行堆叠到几个堆叠的 2D 矩阵中。
具体来说，我对堆叠形状的矩阵很感兴趣 (3,4)其中有 float 元素。例如，如果有 2 个矩阵要堆叠:

import numpy as np
np.arange(24.).reshape(2,3,4)

array([[[ 0.,  1.,  2.,  3.],
        [ 4.,  5.,  6.,  7.],
        [ 8.,  9., 10., 11.]],

       [[12., 13., 14., 15.],
        [16., 17., 18., 19.],
        [20., 21., 22., 23.]]])

并且有一排形状 (1,1,4)堆叠:

row = np.array([[[101.,102.,103.,104.]]])

最终结果将如下所示(堆叠 (4,4) 矩阵):

array([[[  0.,   1.,   2.,   3.],
        [  4.,   5.,   6.,   7.],
        [  8.,   9.,  10.,  11.],
        [101., 102., 103., 104.]],

       [[ 12.,  13.,  14.,  15.],
        [ 16.,  17.,  18.,  19.],
        [ 20.,  21.,  22.,  23.],
        [101., 102., 103., 104.]]])

直到知道，我所做的最好的尝试是使用 np.tile :

import numpy as np
M_stacked = np.arange(24.).reshape(2,3,4)
row = np.array([[[101.,102.,103.,104.]]])
np.concatenate((M_stacked, np.tile(row, (len(M_stacked),1,1))), axis=1)

但我相信这可能不是最有效的解决方案，特别是当堆叠矩阵的数量增加时。有没有更好的方法？
提前致谢!

作为引用，这些是我得到的时间:
如果有 2 个堆叠矩阵:

M_stacked = np.arange(2.*3.*4.).reshape(2,3,4)
%timeit np.concatenate((M_stacked, np.tile(row, (len(M_stacked),1,1))), axis=1)
7.2 µs ± 85 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each))

如果有 1000 个堆叠矩阵:

M_stacked = np.arange(1000.*3.*4.).reshape(1000,3,4)
%timeit np.concatenate((M_stacked, np.tile(row, (len(M_stacked),1,1))), axis=1)
28.8 µs ± 108 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 1

要检查操作的好坏，您可以执行分析内存吞吐量分析。 1 个矩阵大小 (1000, 3, 4)需要读取和 1 个矩阵大小 (1000, 4, 4)需要写。 double 值的大小为 8 个字节(在标准 IEEE-754 兼容系统上)。由于计算在 28.8 us 内完成，因此内存吞吐量为 8*(1000*3*4 + 1000*4*4) / 28.8e-6 / 1024**3 = 7.2 GiB/s这是相对较好但不是很好。
由于计算速度非常快，花费的大部分时间只是 开销调用 Numpy 函数、执行内部数组检查、分配临时 Python 对象(例如元组、Numpy View 、Python 整数)等。
您可以使用 减少开销Numba 的 JIT 仔细调整配置。 Numba 将生成一个函数，该函数将执行更少的检查/分配，而只执行一次。
一种可以大大加快代码速度的解决方案是 手动修复堆叠二维数组的大小 与 Numba。事实上，在非常小的数组上处理可变大小的数组要昂贵得多，因为 Numpy/Numba 不需要使用低效循环迭代最后两个小维度。通过手动修复大小，Numba 可以展开循环并生成更高效的代码(例如，使用 SIMD 指令)。告诉 Numba 矩阵大小的一种简单而优雅的方法是使用断言。如果固定大小不正确，断言对于防止执行错误代码也很有用。
这是结果代码:

import numpy as np
import numba as nb

@nb.njit('f8[:,:,::1](f8[:,:,::1], f8[:,:,::1])')
def compute(M_stacked, row):
    m, n, o = M_stacked.shape
    assert o == 4 # The assertion help numba to generate a much faster code
    assert n == 3
    assert row.shape == (1, 1, o)
    res = np.empty((m, n+1, o), dtype=np.float64)
    for i in range(m):
        # With the above assertions, the following loops will 
        # be unrolled by Numba to produce a very fast code.
        for j in range(n):
            for k in range(o):
                res[i, j, k] = M_stacked[i, j, k]
        for k in range(o):
            res[i, n, k] = row[0, 0, k]
    return res

row = np.array([[[101.,102.,103.,104.]]])
M_stacked = np.arange(1000.*3.*4.).reshape(1000,3,4)
%timeit compute(M_stacked, row)

这是我机器上的结果:

Numpy reference code:           19.5 us (10.7 GiB/s)
Numba code without assertions:  14.9 us (14.0 GiB/s)
Numba code with the assertions:  5.7 us (36.6 GiB/s)

最后一个实现是 3.4 更快 然后是初步实现。对于在 Python 解释器中执行的顺序代码，吞吐量非常好。
请注意，尽管矩阵可能存储在具有更高带宽的 CPU 缓存中，但我的 RAM 的带宽约为 40 GiB/s。此外，我的机器上连续 Numpy 代码(工作 CPU 缓存)的最大实际吞吐量为 56 GiB/s。