Pytorch Inner Product of 3D tensor with 1D Tensor to generate 2D Tensor
Operation : I have pytorch tensor A of dimension [n x m x c] and B of dimension [1 x 1 x c]. I want to take inner product of each of 1 x 1 x c vector from A with B and hence generate a tensor C of dimension [n x m].

Inside forward function of my network at a specific step I receive tensor of dimension [N, channels, Height, Width] where N is number of images, channels is number of channels in the feature map, and height and width are of the current feature map. I also have an [N x channels] feature map from some other subnetwork. In the next step I want to carry out the above mentioned operation.

Can somebody explain the best way and functions available in pytorch to achieve such step.

I am new to pytorch and was unable to find a proper way. Tensorflow supports NHWC format but I think pytorch doesn't, so one of the way is to reshape it to [N, Height, Width, channels] and then iterate like :

# if img is reshaped to [N, H, W, C]
img
# tensor of dimension [N, C]
aud

ans = torch.empty(N, H, W, dtype=torch.double)

for batches in range(img.shape[0]):
    for i in range(img.shape[1]):
        for j in range(img.shape[2]):
            ans[batches][i][j] = torch.dot(img[batches][i][j], aud[batches])


Any other cleaner API ?

PS : This step is required in DeepMind's paper "Object That Sound" for sound localization step.

There is a one-liner

 ans = torch.einsum('nhwc,nc-&gt;nhw', img, aud)


The API of torch.einsum can be difficult to grasp if you haven't had any experience with it before, but it's extremely powerful and generalizes a great deal of liner algebra operations (transpositions, matrix multiplications and traces).

import torch

N, H, W, C = 10, 11, 12, 13
img = torch.randn(N, H, W, C)
aud = torch.randn(N, C)

ans = torch.empty(N, H, W)
for batches in range(img.shape[0]):
    for i in range(img.shape[1]):
        for j in range(img.shape[2]):
            ans[batches][i][j] = torch.dot(img[batches][i][j], aud[batches])

ans2 = torch.einsum('nhwc,nc-&gt;nhw', img, aud)

assert torch.allclose(ans, ans2, atol=1e-6)


Note I had to increase the assertion tolerance above the standard 1e-8 because of numerical precision issues. If einsum becomes a bottleneck in more advanced usecases have a look at opt_einsum which optimizes the order of underlying operations for performance.
