Can't access individual Modules/layers of a Sequential network

neel_g: <p>Hey <img src="https://discuss.pytorch.org/images/emoji/apple/wave.png?v=12" title=":wave:" class="emoji" alt=":wave:" loading="lazy" width="20" height="20"> ,</p>
<p>I had created a <code>Unet</code> with custom encoder from <code>segmentation-models-pytorch</code>. That custom encoder consisted of <code>Module</code>s of its own (which was moved from another network). The result is something like,</p>
<pre data-code-wrap="py"><code class="lang-nohighlight">Unet(
  (encoder): Comma_Encoder()
  (decoder): UnetDecoder(
    (center): Identity()
    (blocks): ModuleList(
      (0): DecoderBlock(
        (conv1): Conv2dReLU(
          (0): Conv2d(131, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (attention1): Attention(
          (attention): Identity()
        )
        (conv2): Conv2dReLU(
          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (attention2): Attention(
          (attention): Identity()
        )
      )
      (1): DecoderBlock(
        (conv1): Conv2dReLU(
          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (attention1): Attention(
          (attention): Identity()
        )
        (conv2): Conv2dReLU(
          (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): ReLU(inplace=True)
        )
        (attention2): Attention(
          (attention): Identity()
        )
      )
    )
  )
  (segmentation_head): SegmentationHead(
    (0): Conv2d(32, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ConvTranspose2d(256, 6, kernel_size=(4, 4), stride=(4, 4))
    (2): Activation(
      (activation): Identity()
    )
  )
)
</code></pre>
<p>However, if I try to access <code>layer</code>s of the <code>encoder</code> by performing <code>model.encoder</code>, I am getting nothing nothing <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<pre data-code-wrap="py"><code class="lang-nohighlight">next(model.encoder.named_modules())
# ('', Comma_Encoder())

[_ for _ in model.encoder.children()]
# []
</code></pre>
<p>Whic is pretty weird since I can forward pass through it easily, and train it correctly too <img src="https://discuss.pytorch.org/images/emoji/apple/person_shrugging.png?v=12" title=":person_shrugging:" class="emoji" alt=":person_shrugging:" loading="lazy" width="20" height="20"> which pretty much implied I just can’t access its submodules.</p>
<p>This may be something with the <code>Unet</code> class, but its inherited from <code>nn.Module</code> too AFAIK (<a href="https://github.com/qubvel/segmentation_models.pytorch/blob/8bf52c7e862af006e76a23aae6aa17977d7f9a79/segmentation_models_pytorch/base/model.py#L5-L6" rel="noopener nofollow ugc">demonstrated here</a>) - and I was able to access the layers of the encoder perfectly before plugging it into SMP…</p>
<p>Is there some sort of limit to the recurrence of some Module? I wanted to utilize accessing individual layers for skip connections which are vital for segmentation <img src="https://discuss.pytorch.org/images/emoji/apple/disappointed.png?v=12" title=":disappointed:" class="emoji" alt=":disappointed:" loading="lazy" width="20" height="20"></p>
<p>Any help is appreciated… <img src="https://discuss.pytorch.org/images/emoji/apple/heart.png?v=12" title=":heart:" class="emoji" alt=":heart:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>I guess <code>Comma_Encoder</code> just doesn’t contain any submodules.<br>
Could you post the class definition here as I cannot find it in the linked repository?</p>

neel_g: <p>sure - the encoder is extracted from an encoder-decoder style architecture (specifically, VQ-VAE-2) as demonstrated in this <a href="https://github.com/rosinality/vq-vae-2-pytorch/blob/ef5f67c46f93624163776caec9e0d95063910eca/vqvae.py#L80-L95" rel="noopener nofollow ugc">repo</a>.</p>
<p>It’s just a bunch of 20 such <code>ResBlock</code>s connected together in 2 ‘submodules’ (<code>encoder_1</code>, and <code>encoder_2</code>)  = 40 blocks connected in series. Everything inherits from <code>Module</code> so I wouldn’t think they wouldn’t be accessible <img src="https://discuss.pytorch.org/images/emoji/apple/person_shrugging.png?v=12" title=":person_shrugging:" class="emoji" alt=":person_shrugging:" loading="lazy" width="20" height="20"></p>
<pre data-code-wrap="py"><code class="lang-nohighlight">class Encoder(nn.Module):
    def __init__(self, in_channel, channel, n_res_block, n_res_channel, stride):
        super().__init__()

        if stride == 4:
            blocks = [
                nn.Conv2d(in_channel, channel // 2, 4, stride=2, padding=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(channel // 2, channel, 4, stride=2, padding=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(channel, channel, 3, padding=1),
            ]

        elif stride == 2:
            blocks = [
                nn.Conv2d(in_channel, channel // 2, 4, stride=2, padding=1),
                nn.ReLU(inplace=True),
                nn.Conv2d(channel // 2, channel, 3, padding=1),
            ]

        for i in range(n_res_block):
            blocks.append(ResBlock(channel, n_res_channel))

        blocks.append(nn.ReLU(inplace=True))

        self.blocks = nn.Sequential(*blocks)

    def forward(self, input):
        return self.blocks(input)
</code></pre>
<p>^ Just basically 2 of these in series make up <code>Comma_Encoder</code>.</p>
<p>Additionally, I am able to view <code>Comma_Encoder</code>'s module when I print it out (before plugging into the Unet) and get exactly what I expected.</p>
<p>I think I may have a clue here - this is how I extract the first few layers from my base model to construct the <code>Encoder</code>:</p>
<pre data-code-wrap="py"><code class="lang-nohighlight">VQVAE = import_mod.VQVAE(**self.base_args).to(self.device)
VQVAE.load_state_dict(self.state_dict)

newmodel = torch.nn.Sequential(*(list(VQVAE.children())))

return [*newmodel[:2].to(self.device), torch.nn.Conv2d(128, 3, 2, 2).to(self.device)]
</code></pre>
<p>It just occured to me - is that a problem due to instantiating <code>Sequential</code>? and in that case, is there any workaround for me to access individual modules while still performing surgery on the model…</p>
<p>Thanks a ton for your immediate response <img src="https://discuss.pytorch.org/images/emoji/apple/rocket.png?v=12" title=":rocket:" class="emoji" alt=":rocket:" loading="lazy" width="20" height="20"> Pytorch forums is really a better place with all its active contributors so willing to help <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>Thanks for the update.<br>
The issue would come from returning a plain Python <code>list</code> containing the modules, which won’t register them properly in the parent module.<br>
Use <code>nn.ModuleList</code> instead:</p>
<pre><code class="lang-python">return nn.ModuleList([*newmodel[:2].to(self.device), torch.nn.Conv2d(128, 3, 2, 2).to(self.device)])
</code></pre>

neel_g: <p>that still doesn’t seem to fix it unfortunately, when iterating over the <code>encoder</code>s <code>named_modules()</code> I still get this</p>
<pre data-code-wrap="py"><code class="lang-nohighlight">[('', Comma_Encoder())]
</code></pre>
<p>which is weirdly nested - as soon as I iterate <code>named_modules</code> in <code>Comma_Encoder()</code> I get the same thing again <img src="https://discuss.pytorch.org/images/emoji/apple/man_shrugging.png?v=12" title=":man_shrugging:" class="emoji" alt=":man_shrugging:" loading="lazy" width="20" height="20"> weird… <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=12" title=":thinking:" class="emoji" alt=":thinking:" loading="lazy" width="20" height="20"></p>

ptrblck: <p>Could you post a minimal, executable code snippet to reproduce the issue, please?</p>

neel_g: <p>Sorry for the late reply, had gone for a vacay…</p>
<p>I was messing around to start some debugging, but apparently If I simple initialize the basemodel in the <code>__init__</code>, everything works as intended:</p>
<pre data-code-wrap="py"><code class="lang-nohighlight">        ....
        self.VQVAE = import_mod.VQVAE(**self.base_args).to(self.device)

    def get_stages(self):
        self.VQVAE.load_state_dict(self.state_dict)
        self.newmodel = torch.nn.Sequential(*(list(VQVAE.children())))
        
        return torch.nn.ModuleList([*newmodel[:2].to(self.device), torch.nn.Conv2d(128, 3, 2, 2).to(self.device)])
</code></pre>
<p>before, everything was under <code>get_stages()</code> to keep things clean.</p>
<p>I’ve no idea how that works, but perhaps you could hazard a guess why such behavior happens for anyone else having trouble?</p>

ptrblck: <p>Both approaches work for me:</p>
<pre><code class="lang-python">class MyModel1(nn.Module):
    def __init__(self):
        super().__init__()
        self.module1 = nn.Linear(1, 1)
        
class MyModel2(nn.Module):
    def __init__(self):
        super().__init__()
        self.get_stages()
        
    def get_stages(self):
        self.module2 = nn.Linear(1, 1)
        
model1 = MyModel1()
print(dict(model1.named_modules()))
# {'': MyModel1(
#   (module1): Linear(in_features=1, out_features=1, bias=True)
# ), 'module1': Linear(in_features=1, out_features=1, bias=True)}

model2 = MyModel2()
print(dict(model2.named_modules()))
# {'': MyModel2(
#   (module2): Linear(in_features=1, out_features=1, bias=True)
# ), 'module2': Linear(in_features=1, out_features=1, bias=True)}
</code></pre>
<p>so I still don’t know what’s causing the issue and can’t add any warnings for it <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>

neel_g: <p>its probably due to the convoluted mess of modules I’ve made <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"> Thanks a lot for your help! &lt;3</p>