這是我的問題。我有兩個矩陣A和B，分別具有復雜的維度(n,n,m,m)和(n,n)。
以下是我為獲得矩陣而執行的操作C-
C = np.sum(B[:,:,None,None]*A, axis=(0,1))
計算上述一次大約需要 6-8 秒。由于我必須計算很多這樣的Cs，所以需要很多時間。有沒有更快的方法來做到這一點？（我在多核 CPU 上使用 JAX NumPy 來做這些；普通的 NumPy 需要更長的時間）
n=77并且m=512，如果您想知道。我可以在處理集群時進行并行化，但陣列的龐大大小會消耗大量記憶體。
uj5u.com熱心網友回復：
看起來你想要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
轉載請註明出處，本文鏈接：https://www.uj5u.com/yidong/384719.html
上一篇：從大型numpy陣列中快速選擇點