python - 如何在 numpy 中执行有效的向量化(内存限制)

Question

我正在处理大量数据，存储在大小为 (3, n) 的 np.array 中。

我需要对每个数据 x: xx.T 执行(以形成 3x3 矩阵)。所以这会产生 n 个 3x3 矩阵。然后，我需要将每个矩阵的结果相加，形成一个 3x3 的整体矩阵。

我按如下方式执行此操作:np.einsum('ij,ki->jk', x, x.T) 这似乎有效。但问题是，我的 n 非常大(我正在处理许多兆像素的非常大的图像)，并且此方法会导致 MemoryError。

在保持矢量化方法的同时，还有其他方法可以做到这一点吗？

谢谢!

Answer 1

根据您的评论进行工作:

In [465]: x=np.array([[1,2,3],[4,5,6]])
In [466]: np.einsum('ij,ik', x,x)
Out[466]: 
array([[17, 22, 27],
       [22, 29, 36],
       [27, 36, 45]])
In [467]: x=np.ones((10,3))
In [468]: np.einsum('ij,ik', x,x)
Out[468]: 
array([[ 10.,  10.,  10.],
       [ 10.,  10.,  10.],
       [ 10.,  10.,  10.]])
In [471]: x=np.ones((10000000,3))
In [472]: np.einsum('ij,ik', x,x)
Out[472]: 
array([[ 10000000.,  10000000.,  10000000.],
       [ 10000000.,  10000000.,  10000000.],
       [ 10000000.,  10000000.,  10000000.]])

如果数组对于内存来说太大，您可以尝试对其进行分块:

In [479]: res = np.zeros((3,3))
In [480]: for i in range(10):
     ...:     res += np.einsum('ij,ik',x,x)
     ...:     
In [481]: res
Out[481]: 
array([[ 170.,  220.,  270.],
       [ 220.,  290.,  360.],
       [ 270.,  360.,  450.]])

通常，对复杂任务进行几次迭代比不进行任何迭代就突破内存限制要快。在某些时候，内存管理成本超过了迭代成本。

<小时/>

这个简单的 einsum 可以用 dot 轻松解决，而且可能更快:

In [484]: x.T.dot(x)
Out[484]: 
array([[17, 22, 27],
       [22, 29, 36],
       [27, 36, 45]])

<小时/>

In [486]: x=np.ones((10000000,3))
In [487]: timeit np.einsum('ij,ik',x,x)
426 ms ± 151 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [488]: x1=np.ones((50000000,3))
In [490]: timeit np.einsum('ij,ik',x1,x1)
2.14 s ± 15.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [493]: %%timeit
     ...: res = np.zeros((3,3))
     ...: x2 = x1.reshape(5,-1,3)
     ...: for i in range(5):
     ...:    x3 = x2[i]
     ...:    res += np.einsum('ij,ik',x3,x3) 
2.1 s ± 5.74 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

因此，即使我还没有达到内存限制，分块也让我的速度略有提高。 (虽然我离内存限制不远了；ipython 历史堆栈中的一些剩余项目可能会把我推倒)。

仅用于比较，@(与 2d 的 dot 相同)

In [494]: %%timeit
     ...: res = np.zeros((3,3))
     ...: x2 = x1.reshape(5,-1,3)
     ...: for i in range(5):
     ...:    x3 = x2[i]
     ...:    res += x3.T@x3
537 ms ± 9.62 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [495]: timeit x1.T@x1
530 ms ± 1.35 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [496]: timeit x.T@x
106 ms ± 1.41 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

在这种大小下，使用点时分块没有帮助。