我需要從頭開始定義矩陣乘法，而不是將每個常數(shù)相乘，每個常數(shù)實際上是另一個數(shù)組，任何兩個數(shù)組都需要“卷積”在一起（我認為沒有必要在這里定義卷積） .我制作了一張圖片，希望能更好地解釋我想說的話：我必須使用的代碼是這樣的：for row in range(arr1.shape[2]): for column in range(arr2.shape[3]): for index in range(arr2.shape[2]): # Could also be "arr1.shape[3]" out[:, :, row, column] += convolve( arr2[:, :, : , column][:, :, index], arr1[:, :, row, : ][:, :, index] )然而，這種方法對我來說非常慢，所以我想知道是否有更快的方法來做到這一點。
1 回答
慕容森
TA貢獻1853條經(jīng)驗 獲得超18個贊
如果中間適合內(nèi)存，則以下內(nèi)容應該相當有效
import numpy as np
from scipy.signal import fftconvolve,convolve
# example
rng = np.random.default_rng()
A = rng.random((5,6,2,3))
B = rng.random((4,3,3,4))
# custom matmul
Ae,Be = A[...,None],B[:,:,None]
shsh = np.maximum(Ae.shape[2:],Be.shape[2:])
Ae = np.broadcast_to(Ae,(*Ae.shape[:2],*shsh))
Be = np.broadcast_to(Be,(*Be.shape[:2],*shsh))
C = fftconvolve(Ae,Be,axes=(0,1),mode='valid').sum(3)
# original loop for reference
out = np.zeros_like(C)
for row in range(A.shape[2]):
for column in range(B.shape[3]):
for index in range(B.shape[2]): # Could also be "A.shape[3]"
out[:, :, row, column] += convolve(
B[:, :, : , column][:, :, index],
A[:, :, row, : ][:, :, index],
mode='valid'
)
print(np.allclose(C,out))
# True
通過批量進行卷積，我們減少了我們必須做的 fft 的總數(shù)。
如果需要，可以通過使用 對傅里葉空間進行總和縮減，進一步優(yōu)化速度和內(nèi)存einsum。不過，這需要手動進行 fft 卷積。
添加回答
舉報
0/150
提交
取消