python - Pytorch 批量矩阵向量外积

Question

我正在尝试使用 PyTorch 生成向量矩阵外积(张量)。假设向量 v 的大小为 p，矩阵 M 的大小为 qXr，则乘积的结果应为pXqXr.

例子:

#size: 2
v = [0, 1] 
#size: 2X3
M = [[0, 1, 2],
     [3, 4, 5]]
#size: 2X2X3
v*M = [[[0, 0, 0],
        [0, 0, 0]],
       [[0, 1, 2],
        [3, 4, 5]]]

对于两个向量 v1 和 v2，我可以使用 torch.bmm(v1.view(1, -1, 1), v2.view( 1, 1, -1))。这可以很容易地扩展到一批向量。但是，我无法找到矢量矩阵案例的解决方案。此外，我需要对批量向量和矩阵执行此操作。

Answer 1

您可以使用 torch.einsum运算符(operator):

torch.einsum('bp,bqr->bpqr', v, M) # batch-wise operation v.shape=(b,p) M.shape=(b,q,r)
torch.einsum('p,qr->pqr', v, M)    # cross-batch operation