Batch-wise and per channe dot prodcut

vision
#1

I have two tensor:
A: [8, 256, 32, 32], which is a feature-map extracted from images
B: [8, 256], which is an intermediate tensor

I want to perform dot product between each channel vector in A and B and output a tensor of size [8, 1, 32, 32]

I use:
torch.matmul(A.permute(0, 3, 2, 1).contiguous().view(-1, 32*32, 256), B.view(-1, 256, 1)).view(-1, 32, 32).unsqueeze(1)

Am I correct? Is there any other straight-forward operation in Pytorch to handle this? Thanks in advance~

#2

What about this?

(A * B.view(B.shape[0], B.shape[1], 1, 1)).sum(dim=1, keepdim=True)
1 Like
#3

I test it and it seems correct. Thanks