How to read this modified unet?
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from PIL import Image
import matplotlib.pyplot as plt

class Model_Down(nn.Module):
    &quot;&quot;&quot;
    Convolutional (Downsampling) Blocks.

    nd = Number of Filters
    kd = Kernel size

    &quot;&quot;&quot;
    def __init__(self,in_channels, nd = 128, kd = 3, padding = 1, stride = 2):
        super(Model_Down,self).__init__()
        self.padder = nn.ReflectionPad2d(padding)
        self.conv1 = nn.Conv2d(in_channels = in_channels, out_channels = nd, kernel_size = kd, stride = stride)
        self.bn1 = nn.BatchNorm2d(nd)

        self.conv2 = nn.Conv2d(in_channels = nd, out_channels = nd, kernel_size = kd, stride = 1)
        self.bn2 = nn.BatchNorm2d(nd)

        self.relu = nn.LeakyReLU()

    def forward(self, x):
        x = self.padder(x)
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.padder(x)
        x = self.conv2(x)
        x = self.bn2(x)
        x = self.relu(x)
        return x

class Model_Skip(nn.Module):
    &quot;&quot;&quot;

    Skip Connections

    ns = Number of filters
    ks = Kernel size

    &quot;&quot;&quot;
    def __init__(self,in_channels = 128, ns = 4, ks = 1, padding = 0, stride = 1):
        super(Model_Skip, self).__init__()
        self.conv = nn.Conv2d(in_channels = in_channels, out_channels = ns, kernel_size = ks, stride = stride, padding = padding)
        self.bn = nn.BatchNorm2d(ns)
        self.relu = nn.LeakyReLU()

    def forward(self,x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        return x


class Model_Up(nn.Module):
    &quot;&quot;&quot;
    Convolutional (Downsampling) Blocks.

    nd = Number of Filters
    kd = Kernel size

    &quot;&quot;&quot;
    def __init__(self, in_channels = 132, nu = 128, ku = 3, padding = 1):
        super(Model_Up, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.padder = nn.ReflectionPad2d(padding)
        self.conv1 = nn.Conv2d(in_channels = in_channels, out_channels = nu, kernel_size = ku, stride = 1, padding = 0)
        self.bn2 = nn.BatchNorm2d(nu)

        self.conv2 =  nn.Conv2d(in_channels = nu, out_channels = nu, kernel_size = 1, stride = 1, padding = 0) #According to supmat.pdf ku = 1 for second layer
        self.bn3 = nn.BatchNorm2d(nu)

        self.relu = nn.LeakyReLU()

    def forward(self,x):
        x = self.bn1(x)
        x = self.padder(x)
        x = self.conv1(x)
        x = self.bn2(x)
        x = self.relu(x)
        x = self.conv2(x)
        x = self.bn3(x)
        x = self.relu(x)
        x = F.interpolate(x, scale_factor = 2, mode = 'bilinear')
        return x


class Model(nn.Module):
    def __init__(self, length = 5, in_channels = 32, out_channels = 3, nu = [128,128,128,128,128] , nd =
                    [128,128,128,128,128], ns = [4,4,4,4,4], ku = [3,3,3,3,3], kd = [3,3,3,3,3], ks = [1,1,1,1,1]):
        super(Model,self).__init__()
        assert length == len(nu), 'Hyperparameters do not match network depth.'

        self.length = length

        self.downs = nn.ModuleList([Model_Down(in_channels = nd[i-1], nd = nd[i], kd = kd[i]) if i != 0 else
                                        Model_Down(in_channels = in_channels, nd = nd[i], kd = kd[i]) for i in range(self.length)])

        self.skips = nn.ModuleList([Model_Skip(in_channels = nd[i], ns = ns[i], ks = ks[i]) for i in range(self.length)])

        self.ups = nn.ModuleList([Model_Up(in_channels = ns[i]+nu[i+1], nu = nu[i], ku = ku[i]) if i != self.length-1 else
                                        Model_Up(in_channels = ns[i], nu = nu[i], ku = ku[i]) for i in range(self.length-1,-1,-1)]) #Elements ordered backwards

        self.conv_out = nn.Conv2d(nu[0],out_channels,1,padding = 0)
        self.sigm = nn.Sigmoid()

    def forward(self,x):
        s = [] #Skip Activations

        #Downpass
        for i in range(self.length):
            x = self.downs[i].forward(x)
            s.append(self.skips[i].forward(x))

        #Uppass
        for i in range(self.length):
            if (i == 0):
                x = self.ups[i].forward(s[-1])
            else:
                x = self.ups[i].forward(torch.cat([x,s[self.length-1-i]],axis = 1))

        x = self.sigm(self.conv_out(x)) #Squash to RGB ([0,1]) format
        return x


This code is a modified UNet I am working on. I am facing hard to read and understand code and how the skip connections are connected to upsampling. can anyone please explain it or can it be written in simple more understandable way without nn.ModuleList.
can some one show how this network looks using a diagram?
This is the github link repo link from where i took this code and trying to understand it.

Here is a functional equivalent of the main Model forward(x) method. It is much more verbose, but it is &quot;unravelling&quot; the flow of operations, making it more easily understandable.
I assumed that the length of the list-arguments are always 5 (i is in the [0, 4] range, inclusive) so I could unpack properly (and it follows the default set of parameters).
def unet_function(x, in_channels = 32, out_channels = 3, nu = [128,128,128,128,128],
                  nd = [128,128,128,128,128], ns = [4,4,4,4,4], ku = [3,3,3,3,3],
                  kd = [3,3,3,3,3], ks = [1,1,1,1,1]):


    ################################
    # DOWN PASS ####################
    ################################

    #########
    # i = 0 #
    #########

    # First Down
    # Model_Down(in_channels = in_channels, nd = nd[i], kd = kd[i])
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2D(in_channels=in_channels, out_channels=nd[0], kernel_size=kd[0], stride=2)(x)
    x = nn.BatchNorm2d(nd[0])(x)
    x = nn.LeakyRelu()(x)
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2d(in_channels = nd[0], out_channels=nd[0], kernel_size = kd[0], stride=1)(x)
    x = nn.BatchNorm2d(nd[0])(x)
    x = nn.LeakyRelu()(x)

    # First skip
    # Model_Skip(in_channels = nd[i], ns = ns[i], ks = ks[i])
    s0 = nn.Conv2D(in_channels=nd[0], out_channels=ns[0])(x)
    s0 = nn.BatchNorm2d(ns[0])(s0)
    s0 = nn.LeakyreLU()(s0)


    #########
    # i = 1 #
    #########

    # Second Down
    # Model_Down(in_channels = nd[i-1], nd = nd[i], kd = kd[i])
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2D(in_channels=nd[0], out_channels=nd[0], kernel_size=kd[1], stride=2)(x)
    x = nn.BatchNorm2d(nd[0])(x)
    x = nn.LeakyRelu()(x)
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2d(in_channels = nd[0], out_channels=nd[0], kernel_size = kd[1], stride=1)(x)
    x = nn.BatchNorm2d(nd[0])(x)
    x = nn.LeakyRelu()(x)

    # Second skip
    # Model_Skip(in_channels = nd[i], ns = ns[i], ks = ks[i])
    s1 = nn.Conv2D(in_channels=nd[1], out_channels=ns[1])(x)
    s1 = nn.BatchNorm2d(ns[1])(s1)
    s1 = nn.LeakyreLU()(s1)


    #########
    # i = 2 #
    #########

    # Third Down
    # Model_Down(in_channels = nd[i-1], nd = nd[i], kd = kd[i])
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2D(in_channels=nd[1], out_channels=nd[1], kernel_size=kd[2], stride=2)(x)
    x = nn.BatchNorm2d(nd[1])(x)
    x = nn.LeakyRelu()(x)
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2d(in_channels = nd[1], out_channels=nd[0], kernel_size = kd[2], stride=1)(x)
    x = nn.BatchNorm2d(nd[1])(x)
    x = nn.LeakyRelu()(x)

    # Third skip
    # Model_Skip(in_channels = nd[i], ns = ns[i], ks = ks[i])
    s2 = nn.Conv2D(in_channels=nd[2], out_channels=ns[2])(x)
    s2 = nn.BatchNorm2d(ns[2])(s2)
    s2 = nn.LeakyreLU()(s2)


    #########
    # i = 3 #
    #########

    # Fourth Down
    # Model_Down(in_channels = nd[i-1], nd = nd[i], kd = kd[i])
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2D(in_channels=nd[2], out_channels=nd[2], kernel_size=kd[3], stride=2)(x)
    x = nn.BatchNorm2d(nd[2])(x)
    x = nn.LeakyRelu()(x)
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2d(in_channels = nd[2], out_channels=nd[2], kernel_size = kd[3], stride=1)(x)
    x = nn.BatchNorm2d(nd[2])(x)
    x = nn.LeakyRelu()(x)

    # Fourth skip
    # Model_Skip(in_channels = nd[i], ns = ns[i], ks = ks[i])
    s3 = nn.Conv2D(in_channels=nd[3], out_channels=ns[3])(x)
    s3 = nn.BatchNorm2d(ns[3])(s3)
    s3 = nn.LeakyreLU()(s3)


    #########
    # i = 4 #
    #########

    # Fifth Down
    # Model_Down(in_channels = nd[i-1], nd = nd[i], kd = kd[i])
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2D(in_channels=nd[3], out_channels=nd[3], kernel_size=kd[4], stride=2)(x)
    x = nn.BatchNorm2d(nd[3])(x)
    x = nn.LeakyRelu()(x)
    x = nn.ReflectionPad2d(padding=1)(x)
    x = nn.Conv2d(in_channels = nd[3], out_channels=nd[3], kernel_size = kd[4], stride=1)(x)
    x = nn.BatchNorm2d(nd[2])(x)
    x = nn.LeakyRelu()(x)

    # Fifth skip
    # Model_Skip(in_channels = nd[i], ns = ns[i], ks = ks[i])
    x = nn.Conv2D(in_channels=nd[4], out_channels=ns[4])(x)
    x = nn.BatchNorm2d(ns[4])(x)
    x = nn.LeakyreLU()(x)



    ################################
    # UP PASS ######################
    ################################

    #########
    # i = 4 #
    #########

    # First Up
    # Model_Up(in_channels = ns[i], nu = nu[i], ku = ku[i])
    x = nn.BatchNorm2d(in_channel=ns[4])(x)
    x = nn.ReflectionPad2d(padding)(x)
    x = nn.Conv2d(in_channels=ns[4], out_channels=nu[4], kernel_size=ku[4], stride=1, padding=0)(x)
    x = nn.BatchNorm2d(nu[4])(x)
    x = nn.LeakyReLU()(x)
    x = nn.Conv2d(in_channels = nu[4], out_channels=nu[4], kernel_size = 1, stride = 1, padding = 0)(x)
    x = nn.BatchNorm2d(nu[4])(x)
    x = nn.LeakyReLU()(x)
    x = F.interpolate(x, scale_factor = 2, mode = 'bilinear')


    #########
    # i = 3 #
    #########

    # Second Up
    # self.ups[i].forward(torch.cat([x,s[self.length-1-i]],axis = 1))
    x = torch.cat([x,s3], axis=1) # IMPORTANT HERE
    # Model_Up(in_channels = ns[i]+nu[i+1], nu = nu[i], ku = ku[i])
    x = nn.BatchNorm2d(in_channel=ns[3]+nu[4])(x)
    x = nn.ReflectionPad2d(padding)(x)
    x = nn.Conv2d(in_channels=ns[3]+nu[4], out_channels=nu[3], kernel_size=ku[3], stride=1, padding=0)(x)
    x = nn.BatchNorm2d(nu[3])(x)
    x = nn.LeakyReLU()(x)
    x = nn.Conv2d(in_channels = ns[3]+nu[4], out_channels=nu[3], kernel_size = 1, stride = 1, padding = 0)(x)
    x = nn.BatchNorm2d(nu[3])(x)
    x = nn.LeakyReLU()(x)
    x = F.interpolate(x, scale_factor = 2, mode = 'bilinear')


    #########
    # i = 2 #
    #########

    # Third Up
    # self.ups[i].forward(torch.cat([x,s[self.length-1-i]],axis = 1))
    x = torch.cat([x,s2], axis=1) # IMPORTANT HERE
    # Model_Up(in_channels = ns[i]+nu[i+1], nu = nu[i], ku = ku[i])
    x = nn.BatchNorm2d(in_channel=ns[2]+nu[3])(x)
    x = nn.ReflectionPad2d(padding)(x)
    x = nn.Conv2d(in_channels=ns[2]+nu[3], out_channels=nu[2], kernel_size=ku[2], stride=1, padding=0)(x)
    x = nn.BatchNorm2d(nu[2])(x)
    x = nn.LeakyReLU()(x)
    x = nn.Conv2d(in_channels = ns[2]+nu[3], out_channels=nu[2], kernel_size = 1, stride = 1, padding = 0)(x)
    x = nn.BatchNorm2d(nu[2])(x)
    x = nn.LeakyReLU()(x)
    x = F.interpolate(x, scale_factor = 2, mode = 'bilinear')


    #########
    # i = 1 #
    #########

    # Fourth Up
    # self.ups[i].forward(torch.cat([x,s[self.length-1-i]],axis = 1))
    x = torch.cat([x,s1], axis=1) # IMPORTANT HERE
    # Model_Up(in_channels = ns[i]+nu[i+1], nu = nu[i], ku = ku[i])
    x = nn.BatchNorm2d(in_channel=ns[1]+nu[2])(x)
    x = nn.ReflectionPad2d(padding)(x)
    x = nn.Conv2d(in_channels=ns[1]+nu[2], out_channels=nu[1], kernel_size=ku[1], stride=1, padding=0)(x)
    x = nn.BatchNorm2d(nu[1])(x)
    x = nn.LeakyReLU()(x)
    x = nn.Conv2d(in_channels = ns[1]+nu[2], out_channels=nu[1], kernel_size = 1, stride = 1, padding = 0)(x)
    x = nn.BatchNorm2d(nu[1])(x)
    x = nn.LeakyReLU()(x)
    x = F.interpolate(x, scale_factor = 2, mode = 'bilinear')    


    #########
    # i = 0 #
    #########

    # Fifth Up
    # self.ups[i].forward(torch.cat([x,s[self.length-1-i]],axis = 1))
    x = torch.cat([x,s0], axis=1) # IMPORTANT HERE
    # Model_Up(in_channels = ns[i]+nu[i+1], nu = nu[i], ku = ku[i])
    x = nn.BatchNorm2d(in_channel=ns[0]+nu[1])(x)
    x = nn.ReflectionPad2d(padding)(x)
    x = nn.Conv2d(in_channels=ns[0]+nu[1], out_channels=nu[0], kernel_size=ku[0], stride=1, padding=0)(x)
    x = nn.BatchNorm2d(nu[0])(x)
    x = nn.LeakyReLU()(x)
    x = nn.Conv2d(in_channels = nu[0], out_channels=nu[0], kernel_size = 1, stride = 1, padding = 0)(x)
    x = nn.BatchNorm2d(nu[0])(x)
    x = nn.LeakyReLU()(x)
    x = F.interpolate(x, scale_factor = 2, mode = 'bilinear')


    ################################
    # OUT ##########################
    ################################

    x = nn.Conv2d(in_channels=nu[0], out_channels=out_channels, kernel_size=1, padding = 0)
    return nn.Sigmoid()(x) #Squash to RGB ([0,1]) format

The two most important parts are:

The skips where the tensor x is processed in a parallel part of the code, not to disturb the main x &quot;pathway&quot;.

The produced tensors from the skip parts are then fed back to the &quot;main pathway&quot; starting from the last one. I kept those tensor as individual variables s0 to s3, so that it is more obvious.


From this picture, you can clearly see the down part feeding latter parts. s0 is the longest grey arrow, it is concatenated to the &quot;main pathway&quot; just before the last convolution layers group.

(Not the same U-Net)
You can also understand from it why we don't need to store a s4 : it is directly fed to the next layer, therefore there is no need to store it as a separate variable.
The Module version does store it, but only because it is conveniently stored in a list read in reverse order at the end. Another obvious reason to store them in a list is so we can have any number of Up and Down parts, by changing the parameters accordingly.
