Pytorch: Memory Efficient weighted sum with weights shared along channels
Inputs:

1) I = Tensor of dim (N, C, X)      (Input)

2) W = Tensor of dim (N, X, Y)    (Weight)

Output:

1) O = Tensor of dim (N, C, Y)     (Output)

I want to compute:

I = I.view(N, C, X, 1)
W = W.view(N, 1, X, Y)
PROD = I*W
O = PROD.sum(dim=2)
return O


without incurring N * C * X * Y memory overhead.

Basically I want to calculate the weighted sum of a feature map wherein the weights are the same along the channel dimension, without incurring memory overhead per channel.



Maybe I could use

from itertools import product

O = torch.zeros(N, C, Y)
for n, x, y in product(range(N), range(X), range(Y)):
    O[n, :, y] += I[n, :, x]*W[n, x, y]
return O


but that would be slower (no broadcasting) and I'm not sure how much memory overhead would be incurred by saving variables for the backward pass.

You can use torch.bmm (https://pytorch.org/docs/stable/torch.html#torch.bmm). Just do torch.bmm(I,W) 

To verify the results : 

import torch
N, C, X, Y= 100, 10, 9, 8 

i = torch.rand(N,C,X)
w = torch.rand(N,X,Y)
o = torch.bmm(i,w)

# desired result code
I = i.view(N, C, X, 1)
W = w.view(N, 1, X, Y)
PROD = I*W
O = PROD.sum(dim=2)

print(torch.allclose(O,o)) # should output True if outputs are same.



EDIT: Ideally, I would assume using pytorch's internal matrix multiplication is efficient. However, you can also measure the memory usage with tracemalloc (at least on CPU). See https://discuss.pytorch.org/t/measuring-peak-memory-usage-tracemalloc-for-pytorch/34067 for GPU. 

import torch
import tracemalloc
tracemalloc.start()
N, C, X, Y= 100, 10, 9, 8 

i = torch.rand(N,C,X)
w = torch.rand(N,X,Y)
o = torch.bmm(i,w)
# output is a tuple indicating current memory and peak memory
print(tracemalloc.get_traced_memory())  


You can do the same with other code and see the bmm implementation is indeed efficient.

import torch
import tracemalloc
tracemalloc.start()
N, C, X, Y= 100, 10, 9, 8 

i = torch.rand(N,C,X)
w = torch.rand(N,X,Y)

I = i.view(N, C, X, 1)
W = w.view(N, 1, X, Y)
PROD = I*W
O = PROD.sum(dim=2)
# output is a tuple indicating current memory and peak memory
print(tracemalloc.get_traced_memory())  


