How to calculate gradients correctly (without in-place operations) for custom unpooling layer?

melike: <p>Hi all, I’m trying to implement a custom max unpooling layer. The difference from <code>torch.nn.MaxUnpool2d</code> is that all the max indices within the pooling window are used for unpooling in case of a repeated max value. Here is an example,</p>
<pre><code class="lang-auto">pooled_inp = torch.tensor([[[[2., 4], 
                         [6, 8]]]], requires_grad=True)

inds = [[[[torch.tensor(3), torch.tensor([0, 1])], 
          [torch.tensor([1, 2, 3]), torch.tensor(0)]]]]

up = MultiMaxUnpool2d(kernel_size=2, stride=2)
unpooled = up(pooled_inp, inds)
print('unpooled:\n', unpooled)

unpooled:
 tensor([[[[0., 0., 4., 4.],
          [0., 2., 0., 0.],
          [0., 6., 8., 0.],
          [6., 6., 0., 0.]]]], grad_fn=&lt;CopySlices&gt;)
</code></pre>
<p>Unpooling works as I expected but the gradients are not calculated correctly. Here are the gradients computed for the above graph,</p>
<pre><code class="lang-auto">unpooled.retain_grad()
unpooled.sum().backward()
print('unpooled.grad:\n', unpooled.grad)
print('pooled_inp.grad:\n', pooled_inp.grad)

unpooled.grad:
 tensor([[[[1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.],
          [1., 1., 1., 1.]]]])

pooled_inp.grad:
 tensor([[[[1., 2.],
          [3., 1.]]]])
</code></pre>
<p>Here is what I expected,</p>
<pre><code class="lang-auto">pooled_inp.grad:
 tensor([[[[1., 1.],
          [1., 1.]]]])
</code></pre>
<p>I don’t understand why the gradients are calculated like that but I’ve learned that the in-place operations should be avoided in Pytorch, so that might be the reason for it. What would be the proper way of implementation without performing in-place operations ? Here is my custom unpooling implementation,</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.nn.modules.utils import _pair

class MultiMaxUnpool2d(nn.Module):
    
    def __init__(self, kernel_size, stride):
        super(MultiMaxUnpool2d, 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)
        
    """ Returns unpooled output size """
    def _calc_size(self, xSize):
        h_in, w_in = xSize[2:]
        h_out = (h_in - 1) * self.stride[0] + self.k[0]
        w_out = (w_in - 1) * self.stride[1] + self.k[1]
        return h_out, w_out
    
    """ Takes an index in 1d, converts it into 2d. Returns its tuple. """
    def _to_2d(self, ind, tl):
        _h, _w = tl
        return (int(ind / self.k[1]) + _h, int(ind % self.k[1]) + _w)
        
    """ Takes a tensor of local indices in a pooling window and that window's
    indices in the pooled input. Returns the list of global indices for 
    unpooled output."""
    def _local_to_global(self, local_inds, i, j):
        tl = (i * self.stride[0], j * self.stride[1])                           # Calculate tl of pooling window in unpooled output. 
        if local_inds.dim() == 0:                                               # torch.tensor() format
            return [self._to_2d(local_inds.item(), tl)]
        else:                                                                   # torch.tensor([...]) format
            return [self._to_2d(ind, tl) for ind in local_inds]
                
    def forward(self, x, inds):
        batch_size, num_ch, h_in, w_in = x.size()[:4]
        h_out, w_out = self._calc_size(x.size())
        unpooled = torch.zeros(batch_size, num_ch, h_out, w_out)
        for b in range(0, batch_size):
            for c in range(0, num_ch):
                for i in range(0, h_in):
                    for j in range(0, w_in):
                        max_val = x[b][c][i][j]
                        local_inds = inds[b][c][i][j]
                        glob_inds = self._local_to_global(local_inds, i, j)     # List of local indices converted to global indices 
                        for g in glob_inds:
                            h_g, w_g = g
                            unpooled[b][c][h_g][w_g] = max_val
        return unpooled
</code></pre>

albanD: <p>Hi,</p>
<p>That gradient actually looks correct to me. These values are used multiple times, so their gradient is the sum of the gradients for each of their contributions.</p>

melike: <p><a class="mention" href="/u/alband">@albanD</a>  In that case, is it proper to use the <code>for</code> loop for indexing the <code>torch.tensor</code> as in the <code>forward</code> function ?</p>
<p>I wrote it before learning that in-place operations should be avoided in PyTorch. And it is stated in <a href="https://discuss.pytorch.org/t/what-is-in-place-operation/16244/15">this</a> reply that the assignment of a value to an index is an in-place operation, so my implementation does not seem correct to me for that reason. But given that the <code>forward</code> and <code>backward</code> functions are working as expected, it might be correct, so I am confused. Any elaboration would be helpful.</p>

albanD: <aside class="quote no-group" data-username="melike" data-post="3" data-topic="105134">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/7ea924/48.png" class="avatar"> melike:</div>
<blockquote>
<p>I wrote it before learning that in-place operations should be avoided in PyTorch.</p>
</blockquote>
</aside>
<p>You don’t have to avoid them. It is just that autograd does not support every combination of them and it will raise an error if you hit such case.<br>
So if your code runs without error, it means that autograd can handle this case just fine.</p>
<p>The only concern I would have with such implementation is the slowdown due to the nested loops. But that’s unrelated to gradient correctness.</p>

melike: <p>Thanks a lot for the clarification, I didn’t know that! The slowdown is not a concern for now but I’ll take a look at it in the future.</p>
<p>All the best.</p>