Pytorch batch matrix vector outer product

问题

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.

Example:

#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]]]

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.

回答1:

you can use 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

回答2:

I was able to do it with following code.

Single vector and matrix

v = torch.arange(3)
M = torch.arange(8).view(2, 4)
# v: tensor([0, 1, 2])
# M: tensor([[0, 1, 2, 3],
#            [4, 5, 6, 7]])

torch.mm(v.unsqueeze(1), M.view(1, 2*4)).view(3,2,4)                                                                                                                                                                              
tensor([[[ 0,  0,  0,  0],
         [ 0,  0,  0,  0]],

        [[ 0,  1,  2,  3],
         [ 4,  5,  6,  7]],

        [[ 0,  2,  4,  6],
         [ 8, 10, 12, 14]]])

For a batch of vectors and matrices, it can be easily extended using torch.bmm.

v = torch.arange(batch_size*2).view(batch_size, 2)
M = torch.arange(batch_size*3*4).view(batch_size, 3, 4)
torch.bmm(v.unsqueeze(2), M.view(-1, 1, 3*4)).view(-1, 2, 3, 4)

回答3:

If [batch_size, z, x, y] is the shape of the target matrix, another solution is building two matrices of this shape with appropriate elements in each position and then apply an elementwise multiplication. It works fine with batch of vectors:

# input matrices
batch_size = 2
x1 = torch.Tensor([0,1])
x2 = torch.Tensor([[0,1,2],
                   [3,4,5]])
x1 = x1.unsqueeze(0).repeat((batch_size, 1))
x2 = x2.unsqueeze(0).repeat((batch_size, 1, 1))

# dimensions
b = x1.shape[0]
z = x1.shape[1]
x = x2.shape[1]
y = x2.shape[2]

# solution
mat1 = x1.reshape(b, z, 1, 1).repeat(1, 1, x, y)
mat2 = x2.reshape(b,1,x,y).repeat(1, z, 1, 1)
mat1*mat2

来源：https://stackoverflow.com/questions/55054127/pytorch-batch-matrix-vector-outer-product

标签

python

pytorch

tensor