How to vectorize local soft-argmax function

Miquel_Espinosa: <p>Hi!</p>
<p>I am trying to implement an efficient parallel and vectorized function to compute the local soft-argmax for a batch of landmarks, where each landmark is a 2D heatmap.</p>
<p>For example, supose I have a tensor of shape: <code>[8,98,128,128]</code>. This would correspond to a tensor containing: <strong>8 batches</strong>, each batch has <strong>98 landmarks</strong>, each landmark contains a <strong>heatmap of 128x128</strong>.</p>
<p>I need to compute the local soft-argmax for each heatmap. The <strong>local soft-argmax</strong> takes the maximum of a heatmap, and computes the soft-argmax locally given some window size.</p>
<p>For example, consider the following scenario.</p>
<ul>
<li>I have a 5x5 example heatmap.</li>
<li>Take a window size of 3x3 around the maximum (<em>in the example is 0.9 at the center</em>)</li>
<li>Compute the soft-argmax around that maximum given the window size</li>
</ul>
<pre><code class="lang-auto"># Original 7x7 heatmap                # Masked heatmap, all zeros except 3x3 around maximum
[0.01, 0.01,  0.01,  0.01, 0.01]      [0.00, 0.00, 0.00, 0.00, 0.00]
[0.01, 0.15,  0.09,  0.03, 0.01]      [0.00, 0.15, 0.09, 0.03, 0.00]
[0.01, 0.80, *0.90*, 0.65, 0.01] --&gt;  [0.00, 0.80, 0.90, 0.65, 0.00] 
[0.01, 0.13,  0.29,  0.33, 0.01]      [0.00, 0.13, 0.29, 0.33, 0.00]
[0.01, 0.01,  0.01,  0.01, 0.01]      [0.00, 0.00, 0.00, 0.00, 0.00]

</code></pre>
<p>And finally, compute the traditional global soft-argmax on the masked heatmap.</p>
<p>I can compute a mask to extract the maximum for each heatmap, but I am unable to efficiently slice its 3x3 neighbours too.</p>
<pre><code class="lang-auto">mask = (output==torch.amax(output, dim=(2,3), keepdim=True))
</code></pre>
<p>Any help would be highly appreciated. Thanks!</p>

Matias_Vasquez: <p>This might help you a little.</p>
<p>Softargmax is missing.</p>
<pre><code class="lang-python">batch_size = 1
nb_channels = 3
h = 10
w = 10
window = 3

conv = torch.nn.Conv2d(nb_channels, nb_channels, window, groups=nb_channels, padding=1, bias=False)
kernel = torch.ones(3, 3).view(1, 1, 3, 3).repeat(nb_channels, 1, 1, 1)
with torch.no_grad():
    conv.weight = torch.nn.Parameter(kernel)


output = torch.rand(batch_size, nb_channels, h, w, requires_grad=True)
mask = (output==torch.amax(output, dim=(2,3), keepdim=True))

masked_output = output.clone()
masked_output[torch.logical_not(mask)] = 0

windowed_masked_output = conv(masked_output)
masked_output[windowed_masked_output&gt;0] = output[windowed_masked_output&gt;0]

torch.set_printoptions(precision=2)
print(output)
print(masked_output)
</code></pre>

Miquel_Espinosa: <p><a class="mention" href="/u/matias_vasquez">@Matias_Vasquez</a>  Wow! Thanks so much!!! This is an incredibly simple yet elegant solution.<br>
I hadn’t thought of using a convolution layer, which of course, makes much sense.</p>
<p>I had a couple of questions regarding the code, so that I fully understand it:</p>
<ul>
<li>
<p>What does exactly do the <code>groups</code> parameter? My intuition is that it separates channels, so that the output is not merged along channels but kept separately. I would appreciate if you could confirm that.<br>
<code>conv = torch.nn.Conv2d(nb_channels, nb_channels, window, groups=nb_channels, padding=1, bias=False)</code></p>
</li>
<li>
<p>Could you further develop on this line? I hadn’t seen this type of indexing before and I struggle to understand the internals.<br>
<code>masked_output[windowed_masked_output&gt;0] = output[windowed_masked_output&gt;0]</code></p>
</li>
</ul>
<hr>
<p>Lastly, to make it work regardless of the window size, I think there is a typo in these two lines, when defining the kernel, and setting padding to ‘same’.</p>
<pre><code class="lang-auto">conv = torch.nn.Conv2d(nb_channels, nb_channels, window, groups=nb_channels, padding='same', bias=False)
kernel = torch.ones(window, window).view(1, 1, window, window).repeat(nb_channels, 1, 1, 1)
</code></pre>
<p>Again, thanks very much for the answer!!</p>

Matias_Vasquez: <aside class="quote no-group" data-username="Miquel_Espinosa" data-post="3" data-topic="150658">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/miquel_espinosa/48/47180_2.png" class="avatar"> Miquel_Espinosa:</div>
<blockquote>
<p>What does exactly do the <code>groups</code> parameter? My intuition is that it separates channels, so that the output is not merged along channels but kept separately. I would appreciate if you could confirm that.<br>
<code>conv = torch.nn.Conv2d(nb_channels, nb_channels, window, groups=nb_channels, padding=1, bias=False)</code></p>
</blockquote>
</aside>
<p>You are correct, this is exactly what is happening.</p>
<aside class="quote no-group" data-username="Miquel_Espinosa" data-post="3" data-topic="150658">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/miquel_espinosa/48/47180_2.png" class="avatar"> Miquel_Espinosa:</div>
<blockquote>
<p>Could you further develop on this line? I hadn’t seen this type of indexing before and I struggle to understand the internals.<br>
<code>masked_output[windowed_masked_output&gt;0] = output[windowed_masked_output&gt;0]</code></p>
</blockquote>
</aside>
<p>When you do a comparison such as <code>==</code>, <code>&lt;</code>, <code>&gt;</code> or others, you get a boolean mask (like you did in your first post).</p>
<p>Now when you use this mask as an index</p>
<pre data-code-wrap="python"><code class="lang-python">masked_output[*MASK*]
</code></pre>
<p>You are only accessing the values where the <code>mask</code>is <code>True</code>. So for your example we would get <code>9</code> values (if the maximum value is <strong>not</strong> on the edge/corner) which share memory with the indexed variable (<code>masked_output</code>in this case). So when we change these values, they change the memory shared, thus changing this variable.</p>
<p>Since <code>output</code>, <code>masked_output</code> and <code>windowed_masked_output</code> have all the same size, there will not be a problem using the same mask everywhere.</p>
<p>So in other words, we select the 9 values that are important to us (which are selected thanks to the <code>windowed_masked_output</code>) from the original <code>output</code> and we copy them into our <code>masked_output</code>, which has just the maximum value and <code>0</code> everywhere else.</p>
<h2><a name="softmax-1" class="anchor" href="#softmax-1"></a>Softmax</h2>
<p>now if you want to do softmax <strong>ONLY</strong> with these 9 values, then you can do something like this.</p>
<pre data-code-wrap="python"><code class="lang-python">m = torch.nn.Softmax(dim=2)
masked_output[windowed_masked_output&gt;0] = m(output[windowed_masked_output&gt;0].view(batch_size, nb_channels, -1)).view_as(output[windowed_masked_output&gt;0])
</code></pre>
<p>which will perform the softmax BEFORE we give it to our <code>masked_output</code></p>
<hr>
<aside class="quote no-group" data-username="Miquel_Espinosa" data-post="3" data-topic="150658">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/miquel_espinosa/48/47180_2.png" class="avatar"> Miquel_Espinosa:</div>
<blockquote>
<p>Lastly, to make it work regardless of the window size, I think there is a typo in these two lines, when defining the kernel, and setting padding to ‘same’.</p>
</blockquote>
</aside>
<p>Thanks, I forgot to change this.</p>

Matias_Vasquez: <p>Oh sorry, there is an error here. This will not work for edge/corner cases.</p>
<ul>
<li>Corner → causes an error.</li>
<li>Edge → wrong answer</li>
</ul>
<aside class="quote no-group" data-username="Matias_Vasquez" data-post="4" data-topic="150658">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/matias_vasquez/48/47847_2.png" class="avatar"> Matias_Vasquez:</div>
<blockquote>
<h2>Softmax</h2>
<p>now if you want to do softmax <strong>ONLY</strong> with these 9 values, then you can do something like this.</p>
<pre><code class="lang-auto">m = torch.nn.Softmax(dim=2)
masked_output[windowed_masked_output&gt;0] = m(output[windowed_masked_output&gt;0].view(batch_size, nb_channels, -1)).view_as(output[windowed_masked_output&gt;0])
</code></pre>
<p>which will perform the softmax BEFORE we give it to our <code>masked_output</code></p>
</blockquote>
</aside>
<h2><a name="this-should-work-better-1" class="anchor" href="#this-should-work-better-1"></a>This should work better</h2>
<pre data-code-wrap="python"><code class="lang-python">m = torch.nn.Softmax(dim=2)
masked_output[windowed_masked_output&gt;0] = output[windowed_masked_output&gt;0]
masked_output[windowed_masked_output==0] = -float('inf')
masked_output = m(masked_output.view(batch_size, nb_channels, -1)).view_as(output)
</code></pre>

Miquel_Espinosa: <aside class="quote no-group" data-username="Matias_Vasquez" data-post="5" data-topic="150658">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/matias_vasquez/48/47847_2.png" class="avatar"> Matias_Vasquez:</div>
<blockquote>
<p><code>masked_output[windowed_masked_output==0] = -float('inf')</code></p>
</blockquote>
</aside>
<p>Thanks for pointing out that detail on the local soft-argmax! I wouldn’t have noticed otherwise, and it is important so that the values outside the window do not count towards that weighted sum.</p>
<p>Thanks again for everything!!!</p>