Why does PyTorch's max pooling layer store input tensors?

username: core-not-dumped
cooked: <p>I made a simple model like below. It seems weird but it has one convolutional layer and two maxpooling layer.</p>
<pre data-code-wrap="python3"><code class="lang-plaintext">class simple_model(nn.Module):
     def __init__(self):
         super(simple_model, self).__init__()
         self.maxpool2D = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
         self.conv1 = nn.Conv2d(3, 20, (5, 5))

     def forward(self, x):
         x = self.maxpool2D(self.maxpool2D(self.conv1(x)))

         return x
</code></pre>
<p>And I check the tensors that saved in forward propagation using gradient hook.</p>
<pre data-code-wrap="python3"><code class="lang-plaintext">pack_saved_tensors = []
def pack_hook(x):
    saved_tensors.append(x)
    return x

unpack_used_tensors = []
def unpack_hook(x):
    unpack_used_tensors.append(x)
    return x

with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
     model_output = model(input_tensors)

label = torch.randn(model_output.size()).to(device)
loss = criterion(model_output, label)
loss.backward()
</code></pre>
<p>This is the result. And I even checked that all stored tensors are used for backward propagation.(with unpack hook funciton) And I think,</p>
<p>0th, 1st tensors saved for convolutional layer,<br>
2nd, 3rd tensors saved for first maxpool2D layer,<br>
4th, 5th tensors saved for second maxpool2D layer.</p>
<pre><code class="lang-auto">pack hook saved tensors:
0 tensor size: torch.Size([64, 3, 224, 224]), tensor type: torch.float32
1 tensor size: torch.Size([20, 3, 5, 5]), tensor type: torch.float32
2 tensor size: torch.Size([64, 20, 220, 220]), tensor type: torch.float32
3 tensor size: torch.Size([64, 20, 110, 110]), tensor type: torch.int64
4 tensor size: torch.Size([64, 20, 110, 110]), tensor type: torch.float32
5 tensor size: torch.Size([64, 20, 55, 55]), tensor type: torch.int64

unpack hook used tensors:
6 tensor size: torch.Size([64, 20, 110, 110]), tensor type: torch.float32
7 tensor size: torch.Size([64, 20, 55, 55]), tensor type: torch.int64
8 tensor size: torch.Size([64, 20, 220, 220]), tensor type: torch.float32
9 tensor size: torch.Size([64, 20, 110, 110]), tensor type: torch.int64
10 tensor size: torch.Size([64, 3, 224, 224]), tensor type: torch.float32
11 tensor size: torch.Size([20, 3, 5, 5]), tensor type: torch.float32
</code></pre>
<p>My question is:</p>
<p>Why Pytorch store input tensors for maxpooling layer? I think, in backward propagation, max pooling layer only need to store int64 tensors(3rd, 5th layer which is store indices of max value).</p>
<p>Any help is appreciated.</p>

username: KFrank
cooked: <p>Hi Core!</p>
<aside class="quote no-group" data-username="core-not-dumped" data-post="1" data-topic="173955" data-full="true">
<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/core-not-dumped/48/57944_2.png" class="avatar"> core-not-dumped:</div>
<blockquote>
<p>Why Pytorch store input tensors for maxpooling layer? I think, in backward propagation, max pooling layer only need to store int64 tensors</p>
</blockquote>
</aside>
<p>My speculation (assuming that built-in pytorch functions work the same<br>
way as do <a href="https://pytorch.org/docs/stable/notes/extending.html#extending-torch-autograd" rel="noopener nofollow ugc">custom autograd functions</a>):</p>
<p>Because <code>MaxPool2D</code> downsamples, it loses some information about the<br>
shape of its input.  When its static <code>backward()</code> method is called, the <em>only</em><br>
thing it knows is <code>grad_output</code> (and hence the shape of its output) and<br>
anything it saved in <code>ctx</code>.  <code>MaxPool2d</code>, as you noted, saves the argmin<br>
index tensor in <code>ctx</code> to avoid recomputation in the backward pass.  But<br>
this index tensor is not enough to reconstruct the shape of the input to<br>
<code>MaxPool2d</code> (which is the same as the shape of the gradient that <code>backward()</code><br>
needs to return).</p>
<p><code>MaxPool2d</code> chooses to save its <code>input</code>, giving it <code>input.shape</code>, which,<br>
together with the argmin index tensor, is enough to compute the gradient.<br>
Saving the input might seem to be overkill, but if its input – often the output<br>
of a previous layer – is already being saved by the previous layer for <em>its</em><br>
backward pass, then saving it again is essentially free.</p>
<p>This is surely not the only way of implementing <code>MaxPool2d</code>, but it seems<br>
economical, so it’s probably a reasonable choice.</p>
<p>Best.</p>
<p>K. Frank</p>

username: core-not-dumped
cooked: <p>Thank you very much for your kind explanation!</p>