Average across last two dimensions in the simplest way

I have tensors of shape N x R x C x H x W (output of ROI pooling). To average across last two dimensions, I currently do:

x.view(x.size()[:3] + (-1,)).mean(-1)

Is there a simpler or more elegant way by using torch.nn.functional.adaptive_avg_pool2d or torch.nn.functional.adaptive_avg_pool3d?

2 Likes

Using view is currently the nicest way, imo.

Looks like someone has been working on something like this in the last week though: https://github.com/pytorch/pytorch/pull/6152

You could also do x.mean(-1).mean(-1). Not terribly pretty, but you don’t need to figure out what the view does… But yes, as Roy mentioned, it might be nice to have multi-dim reductions soon.

Best regards

Thomas

2 Likes