Split network into backbone and head

Mah_Neh: <p>This is my neural network relevant code snippet:</p>
<pre><code class="lang-python">class NeuralNetwork(nn.Module):
    def __init__(self, input_size:int):
        super().__init__()
        self.backbone = SpotRotBackbone()
        # input to the head calculated from the  backbone's structure 
        self.head = SpotRotHead(input_size*input_size*16)
    def forward(self, x):
        x = self.backbone(x)
        return self.head(x)
</code></pre>
<p>Now, the reason why I am using this line:</p>
<p><code>self.head = SpotRotHead(input_size*input_size*16)</code></p>
<p>Is because I know what the backbone does.</p>
<p>Another way I found, much more general but still ugly, is this one:</p>
<pre><code class="lang-python">    def __init__(self, input_size:int):
        super().__init__()
        self.backbone = SpotRotBackbone()
        # can't find a better way
        output_dimensions = self.backbone.forward(torch.zeros(1,3,input_size,input_size)).shape
        self.head = SpotRotHead(output_dimensions[-1])
</code></pre>
<p>What do you think ? How would you do it?</p>

ptrblck: <p>If you cannot create a fixed activation shape after the <code>backbone</code>, e.g. via adaptive pooling layers, both approaches are valid. Pre-computing the input shape or using a fake forward pass are valid approaches. For the latter case you should make sure the actual forward pass using <code>torch.zeros</code> does not change anything inside the <code>self.backbone</code> (e.g. does not update running stats from batchnorm layers etc.).</p>

Mah_Neh: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="186460">
<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>If you cannot create a fixed activation shape after the <code>backbone</code>, e.g. via adaptive pooling layers, both</p>
</blockquote>
</aside>
<p>What do you mean <a class="mention" href="/u/ptrblck">@ptrblck</a> ?</p>
<p>The model’s head looks like this:</p>
<pre><code class="lang-auto">class SpotRotHead(nn.Module):
    """
    Classification of the image rotation angle.
    Args:
        input_size: the size of the input image
    """

    def __init__(self, input_size: int):
        super().__init__()
        self.fc1 = nn.Linear(input_size, 16)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(0.5)
        self.fc2 = nn.Linear(16, 4)

    def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.dropout1(x)
        x = self.fc2(x)
        return x
</code></pre>

ptrblck: <p>Before feeding an activation to the first linear layer, adaptive pooling layers are often used, which allow you to create a defined output shape. If this doesn’t fit your use case, your approaches looks valid.</p>

Mah_Neh: <aside class="quote no-group" data-username="Mah_Neh" data-post="3" data-topic="186460">
<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/mah_neh/48/57251_2.png" class="avatar"> Mah_Neh:</div>
<blockquote>
<p><code>self.fc1 = nn.Linear(input_size, 16)</code></p>
</blockquote>
</aside>
<p>Sounds quite interesting. Do you have any link or code sample so I can see this?</p>

ptrblck: <p>Almost all CNNs use it as seen <a href="https://github.com/pytorch/vision/blob/a7b52a68f625305cfeead988cfd672c167e87a8e/torchvision/models/resnet.py#L278-L280">here</a> for resnet.</p>

Mah_Neh: <p>thank you! a complementary link is<br>
<a href="https://pytorch.org/docs/stable/generated/torch.nn.AdaptiveAvgPool2d.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.nn.AdaptiveAvgPool2d.html</a></p>
<p>So if I get the point this always produce the same output size, which is quite interesting, although I wonder whether using an AdaptiveAvgPool wouldn’t worsen the performance significantly</p>

ptrblck: <p>An aggressive pooling could influence the accuracy and you would need to test it. However, note that it also comes down to allowing users to use a variable input shape vs. raising a shape mismatch error. A reduction in accuracy might thus be the more user-friendly approach instead of a runtime error.</p>