Is there a network layer which requires a 3D input and gives a 2D output?

ivprgrmmr: <p>Hi all,<br>
I am setting up a workflow for a training with DenseNet-121. My data is too small in the depth direction and therefore I have to adjust my network. My input is (288 ×288 ×16) so I want to change the 3D average pooling layer (2x2x2) in my 3rd transition layer to a 2D average pooling layer (2x2).</p>
<p>This is the code I am aiming for in the 3rd transition layer:<br>
(transition3): _Transition(<br>
(norm): BatchNorm3d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)<br>
(relu): ReLU(inplace=True)<br>
(conv): Conv3d(1024, 512, kernel_size=(1, 1, 1), stride=(1, 1, 1), bias=False)<br>
(pool): AvgPool2d(kernel_size=2, stride=2, padding=0)</p>
<p>The problem I am facing is that the input of the AvgPool2D is (18×18×1×512) which is a 3D input, and the AvgPool2D is expecting a 2D input. So I want to ‘squeeze’ (18×18×1×512) to (18×18×512) before using it as input for the pooling layer. However, I do not know how to implement this squeeze function into my transition layer. Does anyone has suggestions on how to fix this problem?</p>
<p>Thanks in advance!</p>

ptrblck: <p>You could write a custom <code>nn.Module</code> and apply the <code>squeeze</code> operation there.<br>
Something like this should work:</p>
<pre><code class="lang-python">class Squeeze(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        
    def forward(self, x):
        x = x.squeeze(dim=self.dim)
        return x

squeeze = Squeeze(dim=2)
x = torch.randn(18, 18, 1, 512)
out = squeeze(x)
print(out.shape)
&gt; torch.Size([18, 18, 512])
</code></pre>

gphilip: <p>How do we figure out if backpropagation can work with such a layer, if we insert it into our network? In this particular, case, <a href="https://pytorch.org/docs/stable/generated/torch.squeeze.html#torch.squeeze" rel="noopener nofollow ugc">the documentation for <code>torch.squeeze()</code></a> doesn’t say whether PyTorch knows how to compute gradients for this operation. Is there a master list of PyTorch functions for which PyTorch knows how to compute gradients?</p>
<p>More generally, how can we ascertain if a custom <code>nn.Module</code> plays well with autograd?</p>

ptrblck: <p>You could run a quick test and check, if the gradients are calculated:</p>
<pre><code class="lang-python">squeeze = Squeeze(dim=2)
x = torch.randn(18, 18, 1, 512, requires_grad=True)
out = squeeze(x)
out.backward(torch.ones_like(out))
print(x.grad)
tensor([[[[1., 1., 1.,  ..., 1., 1., 1.]],

         [[1., 1., 1.,  ..., 1., 1., 1.]],

         [[1., 1., 1.,  ..., 1., 1., 1.]],

         ...,
</code></pre>

gphilip: <p>Thank you.</p>
<p>Just to be clear, this output means that PyTorch does <em>not</em> know how to compute gradients for the Squeeze module, right? Because <code>x.grad</code> remained at its initial all-ones state? Or am I reading this wrong?</p>
<p>Also, what would be a good resource for getting a fairly complete picture of PyTorch autograd? The current picture I have in my head is a mish-mash of graph, leaves, detach, … which is hit-or-miss when it comes to slightly non-trivial scenarios.</p>

ptrblck: <p>No, Autograd knows how to propagate the gradients through <code>squeeze</code>. The ones are expected or which values would you expect when squeezing/unsqueezing a dimension with the size of 1?</p>
<aside class="quote no-group" data-username="gphilip" data-post="5" data-topic="130348">
<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/g/858c86/48.png" class="avatar"> gphilip:</div>
<blockquote>
<p>Also, what would be a good resource for getting a fairly complete picture of PyTorch autograd? The current picture I have in my head is a mish-mash of graph, leaves, detach, … which is hit-or-miss when it comes to slightly non-trivial scenarios.</p>
</blockquote>
</aside>
<p><a href="https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html">This tutorial</a> might be a good starter.</p>

gphilip: <p>Thank you.</p>
<aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="130348">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>No, Autograd knows how to propagate the gradients through <code>squeeze</code>.</p>
</blockquote>
</aside>
<p>What would have been the outcome if Autograd did <em>not</em> know how to do this? An exception/error message saying something about gradients not being computable?</p>
<aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="130348">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>The ones are expected or which values would you expect when squeezing/unsqueezing a dimension with the size of 1?</p>
</blockquote>
</aside>
<p>I cannot figure this out: what is the derivative 
of
 <em>with respect to</em> merely rearranging data? I am not sure if there is a way to define this. But perhaps since there is no  change in the values we can take this derivative to be zero? Is this what Autograd is doing here? It is adding the zeros to the original ones, so the ones stay where they are?</p>
<p>(Please don’t bother answering this question if it is too silly.)</p>
<aside class="quote no-group" data-username="ptrblck" data-post="6" data-topic="130348">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p><a href="https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html" rel="noopener nofollow ugc">This tutorial</a> might be a good starter.</p>
</blockquote>
</aside>
<p>I have skimmed this once or twice. I guess it is time to bite the bullet and go through it more carefully!</p>

ptrblck: <aside class="quote no-group" data-username="gphilip" data-post="7" data-topic="130348">
<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/g/858c86/48.png" class="avatar"> gphilip:</div>
<blockquote>
<p>What would have been the outcome if Autograd did <em>not</em> know how to do this? An exception/error message saying something about gradients not being computable?</p>
</blockquote>
</aside>
<p>Yes, Autograd would yield an error, if the backward method is not available.<br>
This would happen e.g. if you are using a non-differentiable operation:</p>
<pre data-code-wrap="python"><code class="lang-python">x = torch.randn(10, 10, requires_grad=True)
out = torch.argmax(x, dim=1)
out.backward(torch.ones_like(out))
&gt; RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>
<aside class="quote no-group" data-username="gphilip" data-post="7" data-topic="130348">
<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/g/858c86/48.png" class="avatar"> gphilip:</div>
<blockquote>
<p>I cannot figure this out: what is the derivative 
of
 <em>with respect to</em> merely rearranging data? I am not sure if there is a way to define this. But perhaps since there is no change in the values we can take this derivative to be zero? Is this what Autograd is doing here? It is adding the zeros to the original ones, so the ones stay where they are?</p>
</blockquote>
</aside>
<p>No, <code>squeeze</code> and <code>unsqueeze</code> do not change any values and the number of elements of the tensor stays the same so the gradient is just passed to the tensor.</p>

ivprgrmmr: <p>Thanks for the reply. This worked!</p>