numpy vectorized sum over outer products？

luaexgnf 于 2023-10-19 发布在其他

关注(0)|答案(1)|浏览(70)

我在python numpy中有两个float数组：A的形状（60000，）和B的形状（60000，784）。我尝试了以下操作：

np.einsum('i,ijk->jk',A,B[:,:,np.newaxis]*B[:,np.newaxis,:])

但它会产生内存不足错误，因为在中间步骤中创建形状（60000，784，784）的数组需要137GB的RAM。
我尝试了以下替代方法来做到这一点：

S=np.zeros((784,784))
for a,b in izip(A,B):
    S+=a*np.outer(b,b)

但是显式的python sum太慢了。
有没有一种简单的方法来做到这一点，使用矢量化函数来获得快速性能，但在中间步骤中不生成任何（60000，784，784）形状数组？

numpy

来源：https://stackoverflow.com/questions/77045872/numpy-vectorized-sum-over-outer-products

1条答案

按热度按时间

hlswsv351#

看起来你可以简化路径，然后它就不需要那么大的中间数组了。
但是不要忘记optimize=True，否则它会很慢。

import numpy as np

a, b = 60, 784

A = np.random.rand(a)
B = np.random.rand(a, b)

R = np.einsum("i,ijk->jk", A, B[:, :, np.newaxis] * B[:, np.newaxis, :], optimize=True)
R2 = np.einsum("i,ij,ik->jk", A, B, B, optimize=True)

np.allclose(R, R2)  # True

在我的机器上性能很好：

In [1]: a, b = 60_000, 784
   ...: A = np.random.rand(a)
   ...: B = np.random.rand(a, b)
   ...:
   ...: %timeit R = np.einsum("i,ij,ik->jk", A, B, B, optimize=True)
641 ms ± 32.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

赞(0）回复(0）举报 2023-10-19

我来回答

numpy vectorized sum over outer products？

1条答案

相关问题

热门标签

最新问答