我有兩個陣列：
A = torch.rand((64, 128, 10, 10))
B = torch.rand((64, 128, 10))
我想計算由 C 表示的乘積，我們在 A 和 B 的第一維和第二維上進行矩陣向量乘法，所以：
# C should have shape: (64, 128, 10)
for i in range(0, 64):
for j in range(0, 128):
C[i,j] = torch.matmul(A[i,j], B[i,j])
有誰知道如何使用torch.einsum？我嘗試了以下方法，但得到的結果不正確。
C = torch.einsum('ijkl, ijk -> ijk', A, B)
uj5u.com熱心網友回復：
這是帶有 的選項numpy。（我沒有torch）
In [120]: A = np.random.random((64, 128, 10, 10))
...: B = np.random.random((64, 128, 10))
您的迭代參考案例：
In [122]: C = np.zeros((64,128,10))
...: # C should have shape: (64, 128, 10)
...: for i in range(0, 64):
...: for j in range(0, 128):
...: C[i,j] = np.matmul(A[i,j], B[i,j])
...:
matmul完整的廣播：
In [123]: D = np.matmul(A, B[:,:,:,None])
In [125]: C.shape
Out[125]: (64, 128, 10)
In [126]: D.shape # D has an extra size 1 dimension
Out[126]: (64, 128, 10, 1)
In [127]: np.allclose(C,D[...,0]) # or use squeeze
Out[127]: True
einsum等價物：
In [128]: E = np.einsum('ijkl,ijl->ijk', A, B)
In [129]: np.allclose(C,E)
Out[129]: True
轉載請註明出處，本文鏈接：https://www.uj5u.com/ruanti/463179.html
標籤：麻木的 火炬 矩阵乘法 numpy-einsum