我正在寻找一种有效地将同一行堆叠到几个堆叠的2D矩阵的方法。
具体来说，我感兴趣的是具有浮动元素的形状为(3,4)的堆叠矩阵。例如，如果有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)
要检查操作有多好，可以执行分析性内存吞吐量分析。需要读取1个大小为(1000, 3, 4)的矩阵，写入1个大小为(1000, 4, 4)的矩阵。双精度值的大小为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视图，Python整数)等
可以使用Numba的JIT来减少开销通过仔细调优配置。Numba将生成一个函数，它将执行更少的检查/分配，只执行一次。
加快代码速度的一个解决方案是手动固定堆叠的2D数组的大小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解释器中执行的顺序代码来说，吞吐量非常好。
请注意，我的RAM的带宽约为40 GiB/s，尽管矩阵可能存储在具有更高带宽的CPU缓存中。此外，在我的机器上，顺序Numpy代码(工作CPU缓存)的最大实际吞吐量是56 GiB/s。