Pytorch 批量矩阵向量外积 [英] Pytorch batch matrix vector outer product

问题描述

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

I am trying to generate a vector-matrix outer product (tensor) using PyTorch. Assuming the vector v has size p and the matrix M has size qXr, the result of the product should be 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)) .这可以很容易地扩展到一批向量.但是，我无法找到向量矩阵情况的解决方案.另外，我需要对成批的向量和矩阵做这个操作.

For two vectors v1 and v2, I can use torch.bmm(v1.view(1, -1, 1), v2.view(1, 1, -1)). This can be easily extended for a batch of vectors. However, I am not able to find a solution for vector-matrix case. Also, I need to do this operation for batches of vectors and matrices.

推荐答案

你可以使用 einsum

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

这篇关于Pytorch 批量矩阵向量外积的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！