Softargmax in pytorch (in 2D)

Naman-ntc · May 11, 2018, 1:32pm

Softargmax is used quite many place. Deep Spatial Autoencoders for Visuomotor Learning 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!

salihkaragoz · July 14, 2019, 12:35pm

github.com

MWPainter/cvpr2019/blob/master/stitched/soft_argmax.py

from __future__ import division

import torch

def _make_radial_window(width, height, cx, cy, fn, window_width=10.0):
    """
    Returns a grid, where grid[i,j] = fn((i**2 + j**2)**0.5)

    :param width: Width of the grid to return
    :param height: Height of the grid to return
    :param cx: x center
    :param cy: y center
    :param fn: The function to apply
    :return:
    """
    # The length of cx and cy is the number of channels we need
    channels = cx.size(0)

    # Explicitly tile cx and cy, ready for computing the distance matrix below, because pytorch doesn't broadcast very well
    # Make the shape [channels, height, width]

This file has been truncated. show original