[英]Fastest way to multiply and sum 4D array with 2D array in python?
這是我的問題。 我有兩個矩陣A
和B
,具有復雜的條目,尺寸分別為(n,n,m,m)
和(n,n)
。
下面是我執行以獲得矩陣C
的操作 -
C = np.sum(B[:,:,None,None]*A, axis=(0,1))
計算一次以上大約需要 6-8 秒。 由於我必須計算許多這樣的C
,因此需要很多時間。 有沒有更快的方法來做到這一點? (我在多核 CPU 上使用 JAX NumPy 執行這些操作;普通 NumPy 需要更長的時間)
n=77
和m=512
,如果您想知道的話。 當我在集群上工作時,我可以並行化,但是 arrays 的絕對大小消耗了大量的 memory。
看起來你想要einsum
:
C = np.einsum('ijkl,ij->kl', A, B)
在 Colab CPU 上使用 numpy 我得到這個:
import numpy as np
x = np.random.rand(50, 50, 500, 500)
y = np.random.rand(50, 50)
def f1(x, y):
return np.sum(y[:,:,None,None]*x, axis=(0,1))
def f2(x, y):
return np.einsum('ijkl,ij->kl', x, y)
np.testing.assert_allclose(f1(x, y), f2(x, y))
%timeit f1(x, y)
# 1 loop, best of 5: 1.52 s per loop
%timeit f2(x, y)
# 1 loop, best of 5: 620 ms per loop
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.