python - 列 2D 张量乘以行 2D 张量等于 3d pytorch 张量

Question

给定 PyTorch A(a X m) 和(m X b) 中的 2 个二维张量 (a X m) 和B(m X b)，是否有任何有效的方法来获得一个张量C(m X a X b)，其中C[i,:,:] = A[:,i] @ B[i,:]?

在这里，我将举一个问题的例子：

A = torch.FloatTensor([[1,2],[3,4]])
B = torch.FloatTensor([[1,2,3],[4,5,6]])

结果：

C = torch.FloatTensor([[[1,2,3],[3,6,9]],[[12,15,18],[16,20,24]]])

我已经使用for循环完成了它。但是，它的效率非常低。

标签： pythonpytorch

Answer 1

看torch.einsum：

C = torch.einsum('im,mj->mij', A, B)