Almost Convolution Computation

1

Hi there,

I’m new to PyTorch (and deep learning) and I need to do a computation that’s not straight forward (to me) to do with provided functions.

It would be easy to do using “for” but clearly I should avoid that.

The computation is:
For tensor A of dimension (20, 40, 256) and tensor B of dimension (20, 40, 256X300) do an “almost convolution” that results in a tensor C of dimension (20, 40, 300).

Here “almost convolution” means that for each element in the feature map of dim [20, 40] the convolution is done for channels, so the 256 channels from A get filtered through the 300 filters of 256 channels. But at the end, instead of adding all of the results to get a (1,1,300) tensor, the results stay separate for each of the (20,40) elements of the feature maps to get a (20, 40, 300) tensor.

Backpropagation should also work when this computation is part of the computation of the loss function.

2

Hi,

You can try something like this:

A = torch.rand(20, 40, 256)
B = torch.rand(20, 40, 256*300)

C = torch.einsum("ijk, ijkl -> ijl", [A, B.view(20, 40, 256, 300)])

# This is a more generic way
#C = torch.einsum("ijk, ijkl -> ijl", [A, B.view([*A.shape, -1])])

print(C.shape)

# Output: 
torch.Size([20, 40, 300])

Please let me know if this is not the behavior that you expected.

Hope this helps :smile:

1 Like
3

Thank you!

This is the first time I heard of torch.einsum. I worry a bit about how optimal it is, but it should certainly be better than for loops.

I’ll also look into opt_einsum — opt_einsum v3.3.0 documentation

Thank you again.

1 Like