Accessing ModuleList layer

username: saandeep_aathreya
cooked: <p>Below is my rough model structure.</p>
<pre><code class="lang-auto">import torch
from torch import nn
from torch.nn import functional as F
from math import log, pi, exp
import numpy as np
from scipy import linalg as la
from sys import exit as e


class ZeroConv2d(nn.Module):
    def __init__(self, in_channel, out_channel, padding=1):
        super().__init__()

        self.conv = nn.Conv2d(in_channel, out_channel, 3, padding=0)

    def forward(self, input):
        out = self.conv(out)
        return out


class AffineCoupling(nn.Module):
    def __init__(self, in_channel, filter_size=512, affine=True):
        super().__init__()

        self.affine = affine

        self.net = nn.Sequential(
            nn.Conv2d(in_channel // 2, filter_size, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(filter_size, filter_size, 1),
            nn.ReLU(inplace=True),
            ZeroConv2d(filter_size, in_channel if self.affine else in_channel // 2),
        )

        self.net[0].weight.data.normal_(0, 0.05)
        self.net[0].bias.data.zero_()

        self.net[2].weight.data.normal_(0, 0.05)
        self.net[2].bias.data.zero_()

    def forward(self, input):
        in_a, in_b = input.chunk(2, 1)
        log_s, t = self.net(in_a).chunk(2, 1)
        s = F.sigmoid(log_s + 2)
        out_b = (in_b + t) * s
        return torch.cat([in_a, out_b], 1)



class Flow(nn.Module):
    def __init__(self, in_channel, affine=True, conv_lu=True):
        super().__init__()

        self.coupling = AffineCoupling(in_channel, affine=affine)

    def forward(self, input):
        out = self.coupling(out)
        return out


class Block(nn.Module):
    def __init__(self, in_channel, n_flow, split=True, affine=True, conv_lu=True):
        super().__init__()

        squeeze_dim = in_channel * 4

        self.flows = nn.ModuleList()
        for i in range(n_flow):
            self.flows.append(Flow(squeeze_dim, affine=affine, conv_lu=conv_lu))

        self.prior = ZeroConv2d(in_channel * 2, in_channel * 4)


    def forward(self, input):
        for flow in self.flows:
            out = flow(out)
        
        out, z_new = out.chunk(2, 1)
        mean, log_sd = self.prior(out).chunk(2, 1)

        return out, mean, log_sd



class Glow(nn.Module):
    def __init__(self, in_channel, n_flow, n_block):
        super().__init__()

        self.blocks = nn.ModuleList()
        for i in range(n_block - 1):
            self.blocks.append(Block(n_channel, n_flow))
            n_channel *= 2
        self.blocks.append(Block(n_channel, n_flow))

    def forward(self, input):
        out = input
        for block in self.blocks:
            out, mean, log_sd = block(out)
        return out, mean, log_sd

  
</code></pre>
<p>I am trying to access the <code>self.prior = ZeroConv2d()</code> layer under the <code>self.blocks = Block()</code> class. The <code>Block</code> class is a ModuleList defined under the <code>Glow</code> class <code>n_block-1</code> times.</p>
<p>I need to initialize an optimizer where I would like to add all the parameters of this <code>self.prior</code> class.<br>
Right now I am able to index only one ModuleList element to access the <code>self.prior</code> class</p>
<pre><code class="lang-auto">model = Glow()
model.blocks[0].prior
</code></pre>
<p>However, I want something where I can access all the priors but I recieve an error</p>
<pre><code class="lang-auto">model.blocks.prior #ModuleList' object has no attribute 'prior'
</code></pre>
<p>Any help is appreciated.</p>

username: ptrblck
cooked: <p>You could iterate the <code>nn.ModuleList</code> to access each <code>Block</code> and grab its parameters:</p>
<pre><code class="lang-python">model = Glow(1, 1, 10)
params = []
for block in model.blocks:
    params.append(list(block.prior.parameters()))
</code></pre>

username: saandeep_aathreya
cooked: <p>Thank you for your response. How would I add these parameters to my optimizer? Because optim expects my parameters to be of the generator type.</p>
<pre><code class="lang-auto">model = Glow(1, 1, 10)
params = []
for block in model.blocks:
    params.append(list(block.prior.parameters()))

optimizer = optim.Adam(params, lr=lr) #TypeError: optimizer can only optimize Tensors, but one of the params is list
</code></pre>

username: saandeep_aathreya
cooked: <p>Issue is resolved. I had to use extend instead of append.</p>
<pre><code class="lang-auto">model = Glow(1, 1, 10)
params = []
for block in model.blocks:
    params.extend(list(block.prior.parameters()))

optimizer = optim.Adam(params, lr=lr)
</code></pre>