Calculate alpha values with torch mean?

user_14845383 0
Viewed 62

I'm trying to calculate the alpha values as explained here.

I have as argument a tensor with shape (1, 512, 14, 14). To calculate alpha values I need to calculate the average of all dimensions except the channel dimension, so the output will have the shape (1, k, 1, 1) which is essentialy (k,).

How can I do this in PyTorch? Thanks!

pytorch computer-vision
1 Answers

You could permute the first and second axis to keep the channel dimension on dim=0, then flatten all other dimensions, and lastly, take the mean on that new axis:

x.permute(1, 0, 2, 3).flatten(start_dim=1).mean(dim=1)

Here are the shapes, step by step:

>>> x.permute(1, 0, 2, 3).shape
(512, 1, 14, 14)

>>> x.permute(1, 0, 2, 3).flatten(start_dim=1).shape
(512, 1, 196)

>>> x.permute(1, 0, 2, 3).flatten(start_dim=1).mean(dim=1).shape
(512,)
user_63313690
Related