Build blocks inside or outside the network class?

What are the drawbacks/advantages of defining a function that builds a given layer as a member of the neural network class as opposed to defining it outside the class?

For example, as a class member:

class Discriminator(nn.Module):
    """C64-C128-C256-C512 PatchGAN Discriminator architecture."""

    def __init__(self, in_channels=3):
        super(Discriminator, self).__init__()
        channels = [in_channels, 64, 128, 256, 512, 1]
        layers = len(channels)
        channels = zip(channels, channels[1:])

        self.blocks = nn.ModuleList()
        for layer, (input_size, output_size) in enumerate(channels):
            if layer == 0:
                input_size *= 2
                batch_norm = False
                leaky = True
            elif layer < layers - 2:
                batch_norm = True
                leaky = True
            else:
                batch_norm = False
                leaky = False
            self.blocks.append(self.make_conv(in_size=input_size,
                                              out_size=output_size,
                                              batch_norm=batch_norm,
                                              leaky=leaky))

        self.init_weights(mean=0.0, std=0.02)

    def make_conv(self, in_size, out_size, batch_norm, leaky):
        """Convolutional blocks of the Discriminator.

        Let Ck denote a Convolution-BtachNorm-ReLU block with k channels.
        All convolutions are 4 x 4 spatial filters with stride 2 and
        downsample by a factor of 2. BatchNorm is not applied to the c64 block.

        After the C512 block, a convolution is applied to map to a 1-d output,
        followed by a Sigmoid function. All ReLUs are leaky with slope of 0.2.
        """
        block = [nn.Conv2d(in_size, out_size,
                           kernel_size=4, stride=2, padding=2,
                           padding_mode="reflect",
                           bias=False if batch_norm else True)]
        if batch_norm:
            block.append(nn.BatchNorm2d(out_size))
        if leaky:
            block.append(nn.LeakyReLU(0.2))
        else:
            block.append(nn.Sigmoid())

        return nn.Sequential(*block)

    def init_weights(self, mean=0.0, std=0.02):
        """Initialize weights from a Gaussian distribution."""
        for module in self.modules():
            if isinstance(module, (nn.Conv2d, nn.BatchNorm2d)):
                nn.init.normal_(module.weight.data, mean=mean, std=std)

    def forward(self, x, y):
        """Output an nxn tensor belonging to patch ij in the input image."""
        x = torch.cat((x, y), dim=1)
        for block in self.blocks:
            x = block(x)
        return x

outside the class:


def make_conv(self, in_size, out_size, batch_norm, leaky):
   """Convolutional blocks of the Discriminator.

    Let Ck denote a Convolution-BtachNorm-ReLU block with k channels.
    All convolutions are 4 x 4 spatial filters with stride 2 and
    downsample by a factor of 2. BatchNorm is not applied to the c64 block.

    After the C512 block, a convolution is applied to map to a 1-d output,
    followed by a Sigmoid function. All ReLUs are leaky with slope of 0.2.
    """
    block = [nn.Conv2d(in_size, out_size,
                           kernel_size=4, stride=2, padding=2,
                           padding_mode="reflect",
                           bias=False if batch_norm else True)]
    if batch_norm:
        block.append(nn.BatchNorm2d(out_size))
    if leaky:
        block.append(nn.LeakyReLU(0.2))
    else:
        block.append(nn.Sigmoid())

    return nn.Sequential(*block)

class Discriminator(nn.Module):
    """C64-C128-C256-C512 PatchGAN Discriminator architecture."""

    def __init__(self, in_channels=3):
        super(Discriminator, self).__init__()
        channels = [in_channels, 64, 128, 256, 512, 1]
        layers = len(channels)
        channels = zip(channels, channels[1:])

        self.blocks = nn.ModuleList()
        for layer, (input_size, output_size) in enumerate(channels):
            if layer == 0:
                input_size *= 2
                batch_norm = False
                leaky = True
            elif layer < layers - 2:
                batch_norm = True
                leaky = True
            else:
                batch_norm = False
                leaky = False
            self.blocks.append(make_conv(in_size=input_size,
                                              out_size=output_size,
                                              batch_norm=batch_norm,
                                              leaky=leaky))

        self.init_weights(mean=0.0, std=0.02)

    def init_weights(self, mean=0.0, std=0.02):
        """Initialize weights from a Gaussian distribution."""
        for module in self.modules():
            if isinstance(module, (nn.Conv2d, nn.BatchNorm2d)):
                nn.init.normal_(module.weight.data, mean=mean, std=std)

    def forward(self, x, y):
        """Output an nxn tensor belonging to patch ij in the input image."""
        x = torch.cat((x, y), dim=1)
        for block in self.blocks:
            x = block(x)
        return x

I think it mainly boils down to your coding style and in particular if and where you would like to reuse this method.
I don’t see any issues with writing make_conv as a class method in Discriminator in case this is the only model using it. Otherwise, I would define it outside of the class so that other models and scripts could call it too.

1 Like