How to implement dim parameter of torch.max()?

username: melike
cooked: <p>Hi all, I am trying to implement a <a href="https://discuss.pytorch.org/t/implementing-a-variation-of-maxpool2d/102543">variation</a> of <code>torch.max()</code> and confused about the implementation of <code>dim</code> parameter. Here is an example,</p>
<pre><code class="lang-auto">inp = torch.arange(0., 24).view(2, 3, 4)
torch.max(inp, dim=1)

torch.return_types.max(
values=tensor([[ 8.,  9., 10., 11.],
        [20., 21., 22., 23.]]),
indices=tensor([[2, 2, 2, 2],
        [2, 2, 2, 2]]))
</code></pre>
<p>So, the max value within each colored rectangle is returned as in the figure.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/5/b5bf49b1c56347952fbb6637f6e0fe4e7caafda0.jpeg" alt="dim=1" data-base62-sha1="pVOfMngwGwbwLxP5iLkBH7e9beo" width="329" height="133"></p>
<p>Likewise, <code>torch.max(inp, dim=2)</code> returns the max value of each colored rectangle in the below figure,</p>
<pre><code class="lang-auto">torch.return_types.max(
values=tensor([[ 3.,  7., 11.],
        [15., 19., 23.]]),
indices=tensor([[3, 3, 3],
        [3, 3, 3]]))
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/3/d361489a46fc1c536a661f97553715fab2a5d9c4.jpeg" alt="dim=2" data-base62-sha1="u9XaSkKexFLBZkWbfKChEpYUj8E" width="287" height="134"></p>
<p>My question is how can I reshape the input as done with the <code>torch.max()</code> in order to use <code>dim</code> parameter in my custom max implementation ? I’ve experimented <code>view</code> with different values but it didn’t work so far.</p>

username: ptrblck
cooked: <p>What is your custom <code>max</code> operation doing differently and what is not working at the moment?</p>

username: melike
cooked: <p>The custom <code>max</code> should return the indices of all maximum values instead of the first one being encountered as in <code>torch.max</code>. I want to add <code>dim</code> as a parameter to my custom <code>max</code>, like in <code>torch.max</code>. But I don’t know how to implement the dimension reduction effect controlled by <code>dim</code>.</p>
<p><a href="https://pytorch.org/docs/stable/generated/torch.max.html" rel="noopener nofollow ugc">Documentation</a> of <code>torch.max</code> states as below but the input is not in the shape of <code>(Ax1xB)</code>, so using <code>torch.squeeze</code> does not yield to reduced dimensions.</p>
<blockquote>
<p>Otherwise, dim is squeezed (see <code>torch.squeeze()</code>), resulting in the output tensors having 1 fewer dimension than input</p>
</blockquote>

username: ptrblck
cooked: <p>If you want to return more than a single <code>max</code> value, you could try to use e.g. <code>topk</code> instead of <code>max</code> or sort the tensor and grab the values manually.</p>

username: melike
cooked: <p>The <a href="https://pytorch.org/docs/stable/generated/torch.topk.html" rel="noopener nofollow ugc">documentation</a> of <code>topk</code> says it returns the k largest elements but I only need the max values and don’t know how many of them will be there.</p>
<p>I tried using <code>max</code> but it results in only the first max value having the gradient, it’s not shared among the max values. Here is an example,</p>
<pre><code class="lang-auto">inp = torch.tensor([[ 6.,  2,  3,  6],
                    [ 9, 16, 16, 16]], requires_grad=True)
outp, inds = torch.max(inp, dim=-1)
outp.sum().backward()
print('outp:\n', outp)
print('inp.grad:\n', inp.grad)

outp:
 tensor([ 6., 16.], grad_fn=&lt;MaxBackward0&gt;)
inp.grad:
 tensor([[1., 0., 0., 0.],
        [0., 1., 0., 0.]])
</code></pre>
<p>I expect the gradient to be like below,</p>
<pre><code class="lang-auto">inp.grad:
 tensor([[0.5, 0., 0., 0.5],
        [0., 0.333, 0.333, 0.333]])
</code></pre>
<p>So, I think I have to implement a custom <code>max</code> but I don’t know how to reshape the input (or squeeze it) controlled by the <code>dim</code> parameter in <code>torch.max</code>.</p>

username: ptrblck
cooked: <p>Thanks for the follow-up as I misunderstood your use case.<br>
Based on your description it seems that the returned values would have a variable shape. E.g. one window could have the same max values, while another one a single one.<br>
If that’s the case I think you would have to return a <code>list</code> of these values e.g. by getting the max value via <code>torch.max</code> and checking if duplicates can be found in the current window.</p>
<p>I’m not sure how you would like to further process the result <code>list</code>, as you wouldn’t be able to create a single return tensor (with nested tensors this might be possible).</p>

username: melike
cooked: <p>Thank you for bearing with me, this post will be a bit long. The returned values by <code>max</code> don’t need to be variable in shape since they keep the max value and there is only one (maybe repeating) max within the window, but the returned indices will be variable in shape. An example is at the <code>MultiMax</code> part.</p>
<p>I’ve implemented the custom <code>max</code> (named <code>MultiMax</code>) as a custom autograd function and the max pooling (named <code>MultiMaxPool2d</code>) layer using it. As I asked in this thread, I don’t know how to implement the <code>dim</code> parameter of <code>torch.max</code> for my <code>MultiMax</code>. So, <code>MultiMax</code> takes a 1d tensor and is called within the <code>for</code> loop of <code>MultiMaxPool2d</code> layer. I’ve just learned this approach might be problematic since in-place operations should be avoided in PyTorch.</p>
<p>My question is, what should be the proper approach instead of using the in-place operations for creating the pooled output ? I obviously need to select max values from the input and assign them to their new positions in the pooled output. I’ve added the source code and examples.</p>
<p>Here is my custom <code>max</code> operation <code>MultiMax</code>,</p>
<pre><code class="lang-auto">import torch

class MultiMax(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        mmax = torch.max(input)
        inds = torch.nonzero(mmax == input).squeeze() # Get max inds
        ctx.save_for_backward(input, mmax, inds)
        ctx.mark_non_differentiable(inds)
        return mmax, inds
    
    @staticmethod
    def backward(ctx, grad_output, for_inds):
        input, mmax, inds = ctx.saved_tensors
        inds_shape = 1 if inds.dim() == 0 else inds.shape[0]
        grad_input = torch.ones(input.shape[0]) * grad_output
        grad_input /= inds_shape  # Grad is shared among max values. 
        grad_input[input &lt; mmax] = 0
        return grad_input, None
</code></pre>
<p>An example using <code>MultiMax</code>,</p>
<pre><code class="lang-auto">inp = torch.tensor([6., 2, 5, 6], requires_grad=True)
mmax = MultiMax.apply
pooled, inds = mmax(inp)
pooled.sum().backward()
print('pooled:\n', pooled)
print('inds:\n', inds)
print('inp grads:\n', inp.grad)

pooled:
  tensor(6., grad_fn=&lt;MultiMaxBackward&gt;)
inds:
  tensor([0, 3])
inp grads:
  tensor([0.5000, 0.0000, 0.0000, 0.5000])
</code></pre>
<p>Here is the implementation of custom max pooling <code>MultiMaxPool2d</code> using <code>MultiMax</code> (adapted from <a href="https://github.com/ChristoferNal/custom-pooling-operation/blob/master/custom_pooling.ipynb" rel="noopener nofollow ugc">here</a>),</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.modules.utils import _pair, _quadruple
from MultiMax import MultiMax

class MultiMaxPool2d(nn.Module):
    
    def __init__(self, kernel_size, stride, padding=0, same=False):
        super(MultiMaxPool2d, self).__init__()
        self.k = kernel_size if type(kernel_size) is tuple else _pair(kernel_size)
        self.stride = stride if type(stride) is tuple else _pair(stride)
        self.padding = padding if type(padding) is tuple else _quadruple(padding)
        self.same = same
        
    def _init_pool_inds(self, s):
        pool = torch.zeros(s[0], s[1], s[2], s[3])
        inds = []
        for x in range(0, s[0]):
            inds.append([])
            for y in range(0, s[1]):
                inds[x].append([])
                for z in range(0, s[2]):
                    inds[x][y].append([])
                    for t in range(0, s[3]):
                        inds[x][y][z].append(0)
        return pool, inds
    
    def _padding(self, x):
        if self.same:
            ih, iw = x.size()[2:]
            if ih % self.stride[0] == 0:
                ph = max(self.k[0] - self.stride[0], 0)
            else:
                ph = max(self.k[0] - (ih % self.stride[0]), 0)
            if iw % self.stride[1] == 0:
                pw = max(self.k[1] - self.stride[1], 0)
            else:
                pw = max(self.k[1] - (iw % self.stride[1]), 0)
            pl = pw // 2
            pr = pw - pl
            pt = ph // 2
            pb = ph - pt
            padding = (pl, pr, pt, pb)
        else:
            padding = self.padding
        return padding
        
    def forward(self, x):
        x = F.pad(x, self._padding(x), mode='reflect')
        x = x.unfold(2, self.k[0], self.stride[0]).unfold(3, self.k[1], self.stride[1])
        x = x.contiguous().view(x.size()[:4] + (-1,))
        mmax = MultiMax.apply
        s = x.shape[:4]
        pool, inds = self._init_pool_inds(s)
        for i in range(0, s[0]):
            for j in range(0, s[1]):
                for k in range(0, s[2]):
                    for l in range(0, s[3]):
                        _max, _is = mmax(x[i][j][k][l])
                        pool[i][j][k][l] = _max
                        inds[i][j][k][l] = _is
        return pool, inds
</code></pre>
<p>An example using <code>MultiMaxPool2d</code>,</p>
<pre><code class="lang-auto">inp = torch.tensor([[[[ 6.,  2,  3,  9],
                      [ 5,  6,  7,  8],
                      [ 9, 10, 11, 16],
                      [14, 14, 16, 16]]]], requires_grad=True)

mp = MultiMaxPool2d(kernel_size=2, stride=2)
pooled, inds = mp(inp)
pooled.sum().backward()
print('pooled:\n', pooled)
print('inds:\n', inds)
print('inp grads:\n', inp.grad)

pooled:
 tensor([[[[ 6.,  9.],
          [14., 16.]]]], grad_fn=&lt;CopySlices&gt;)
inds:
 [[[[tensor([0, 3]), tensor(1)], [tensor([2, 3]), tensor([1, 2, 3])]]]]

inp grads:
 tensor([[[[0.5000, 0.0000, 0.0000, 1.0000],
          [0.0000, 0.5000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.3333],
          [0.5000, 0.5000, 0.3333, 0.3333]]]])
</code></pre>
<p>The gradients are calculated as I expected in both of the operations but I also implemented an unpooling layer which I believe also suffers from in-place operations and the gradients are not calculated as I expected.</p>

username: melike
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I still don’t know how to implement the dimension reduction for my custom <code>max</code> operation as controlled by the <code>dim</code> parameter of <code>torch.max</code> but this <a href="https://discuss.pytorch.org/t/how-to-calculate-gradients-correctly-without-in-place-operations-for-custom-unpooling-layer/105134/4">reply</a> made it clear that if the code runs without errors, then autograd handles the in-place operations.</p>