Accessing ModuleList layer

vision
1

Below is my rough model structure.

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

I am trying to access the self.prior = ZeroConv2d() layer under the self.blocks = Block() class. The Block class is a ModuleList defined under the Glow class n_block-1 times.

I need to initialize an optimizer where I would like to add all the parameters of this self.prior class.
Right now I am able to index only one ModuleList element to access the self.prior class

model = Glow()
model.blocks[0].prior

However, I want something where I can access all the priors but I recieve an error

model.blocks.prior #ModuleList' object has no attribute 'prior'

Any help is appreciated.

2

You could iterate the nn.ModuleList to access each Block and grab its parameters:

model = Glow(1, 1, 10)
params = []
for block in model.blocks:
    params.append(list(block.prior.parameters()))
1 Like
3

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.

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
4

Issue is resolved. I had to use extend instead of append.

model = Glow(1, 1, 10)
params = []
for block in model.blocks:
    params.extend(list(block.prior.parameters()))

optimizer = optim.Adam(params, lr=lr)