Softargmax is used quite many place. Deep Spatial Autoencoders for Visuomotor Learning 99 probably introduced it.
Is there any function/layer in pytorch that performs it or any custom implementation.
I implemented it on following idea :
# consider input [N,C,H,W] (batch, channels, height, width)
temp = input.view(N,C,-1)
m = torch.softmax(dim=-1)
weights = m(temp)
semi_indices = (weights*(torch.arange(H*W).unsqueeze(0).unsqueeze(0).expand(weights.size()))).sum(dim=-1)
indices_x = semi_indices / H
indices_y = semi_indices % W
# maths could be off a bit but basically somewhat like this

But I really didnt like this implementation, please suggest something (maybe using Softmax2d but it similar i guess). Would there be any other way to do this please suggest!