Wrapping submodules in Sequential

username: seer_mer
cooked: <p>Is the statement in the link below true?</p><aside class="quote quote-modified" data-post="27" data-topic="17648">
  <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">
    <a href="https://discuss.pytorch.org/t/how-to-delete-layer-in-pretrained-model/17648/27">How to delete layer in pretrained model?</a> 
  </div>
  <blockquote>
    InceptionV3 uses functional calls in its forward, which will be missing, if you wrap all submodules in an nn.Sequential container. 
In particular <a href="https://github.com/pytorch/vision/blob/dec8628d92c52e9b0d3f8816037870907b7d99f2/torchvision/models/inception.py#L174" rel="noopener nofollow ugc">this flatten operation</a> will be missing. 

You could use forward hooks as described <a href="https://discuss.pytorch.org/t/how-can-l-load-my-best-model-as-a-feature-extractor-evaluator/17254/6">here</a>.
  </blockquote>
</aside>

<p>According to the response, functional calls will be missing when using nn.Sequential.<br>
However, I saw many repos use submodules with functional api in sequential, for example, the official convnext implementation:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py">
  <header class="source">

      <a href="https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py" target="_blank" rel="noopener nofollow ugc">facebookresearch/ConvNeXt/blob/main/models/convnext.py</a></h4>


      <pre><code class="lang-py"># Copyright (c) Meta Platforms, Inc. and affiliates.

# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.


import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import trunc_normal_, DropPath
from timm.models.registry import register_model

class Block(nn.Module):
    r""" ConvNeXt Block. There are two equivalent implementations:
    (1) DwConv -&gt; LayerNorm (channels_first) -&gt; 1x1 Conv -&gt; GELU -&gt; 1x1 Conv; all in (N, C, H, W)
    (2) DwConv -&gt; Permute to (N, H, W, C); LayerNorm (channels_last) -&gt; Linear -&gt; GELU -&gt; Linear; Permute back
    We use (2) as we find it slightly faster in PyTorch
    
</code></pre>



  This file has been truncated. <a href="https://github.com/facebookresearch/ConvNeXt/blob/main/models/convnext.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>
<p>
They used functional api (F.layer_norm) inside submodule (class Block), and used nn.Sequential to wrap all the blocks.</p>

username: Matias_Vasquez
cooked: <p>Maybe the easiest way to understand this is with a little example.</p>
<p>Here I have defined two Modules that perform functional stuff. The <code>SimpleBlock</code> multiplies by <code>2</code>. The <code>MainModel</code> uses this <code>SimpleBlock</code> in two different places and also performs some functional stuff (multiply by <code>3</code> and add <code>3</code>).</p>
<p>This can, of course, be something more complex like using layers, batch normalization, flatten or whatever you want to do.</p>
<pre><code class="lang-python">class SimpleBlock(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def functional_stuff(self, x):
        return x * 2
    
    def forward(self, x):
        # Functional - Multiply by 2
        x = self.functional_stuff(x)
        return x

class MainModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.block_0 = SimpleBlock()
        self.block_1 = SimpleBlock()
        
    def functional_stuff(self, x):
        # This is inside a nn.Module
        x = self.block_0(x)

        # This is functional
        x = x * 3
        return x

    def forward(self, x):
        x = self.functional_stuff(x)

        # Functional - Add 3
        x = x + 3 

        # This is a nn.Module
        # The functional stuff indside this module is done
        # even if you extract the children with Sequential
        x = self.block_1(x)
        return x
</code></pre>
<p>We now use the <code>MainModel</code> and extract the children using <code>nn.Sequential</code>.</p>
<pre><code class="lang-python">input = torch.ones(2, 2)

model = MainModel()
seq_model = torch.nn.Sequential(*list(model.children()))

print(model(input))
print(seq_model(input))
</code></pre>
<p>If we now see the output generated, we can see that they do NOT yield the same result.</p>
<pre><code class="lang-python"># Output

# Model
tensor([[18., 18.],
        [18., 18.]])
# Sequential extraction
tensor([[4., 4.],
        [4., 4.]])
</code></pre>
<p>The full model gives the correct output.<br>
<code>([1 * 2 ] * 3 + 3) * 2 = ([2] * 3 + 3) * 2 = (6 + 3)*2 = (9)*2=18</code></p>
<p>However the <code>nn.Sequential</code> did <strong>not</strong> multiply by <code>3</code> or added the <code>3</code>.<br>
<code>([1*2])*2 = 2*2 = 4</code></p>
<p>This is because the functional stuff is wrapped inside this <code>nn.Module</code> called <code>SimpleBlock</code>. But the stuff on the <code>forward</code> method of our <code>MainModel</code> is not. We can even use <code>nn.Modules</code> that are in another function other than <code>forward</code> and call them from inside <code>forward</code>. But if they are <strong>not</strong> <code>nn.Modules</code> (or a container like <code>nn.Sqeuential</code> or others) they will not work.</p>
<p>In your example of <code>ConvNeXt</code>, this <code>F.layer_norm</code> is inside the class <code>LayerNorm</code> used there. This will work.</p>
<p>Also in the forward method of <code>ConvNeXt</code> we call another method. Inside this method. However, everything here is either a <code>nn.Module</code> or inside a container like <code>nn.Sequential</code> and <code>nn.ModuleList</code>.</p>
<p>So if we did something like this</p>
<pre><code class="lang-python">mod = torchvision.models.convnext_tiny(pretrained=True)
seq_mod = torch.nn.Sequential(*list(mod.children()))

mod.eval()
seq_mod.eval()

input = torch.rand(1, 3, 300, 300)

print(torch.all(mod(input) == seq_mod(input)))
</code></pre>
<p>We get <code>tensor(True)</code> as the output.<br>
If there was any other form of functional API on the forward of <code>ConvNeXt</code> (not the children but ONLY on <code>ConvNeXt</code>) then it would not work the same.</p>
<p>In the example given, there is a <code>flatten</code> function, which will break the pipeline if missing.</p>
<p>Hope this helps and makes it a bit clearer <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>