What's the correct way to implement convolutional blocks of specified depth?

Question

I'm trying to implement bayesian optimization with Convolutional Neural Netowrk in PyTorch, specifically, I'm trying to translate network structure from Matlab BayesOptExperiment to PyTorch. I want my network to have the following structure:

Input data -> convblock -> maxpool -> convblock -> maxpool -> convblock -> avgpool -> flatten -> linear -> softmax

where convblock consists of:

[conv2Dlayer -> batch normalization layer -> ReLU],

repeated a few times. Current version works as expected only when section_depth = 1 achieving accuracy of around 65-70%, although if I raise the depth of convblock accuracy plummets to around 10%. I'm definitely missing something, but I'm not sure what is that. The structure of my network:

import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

#...

class Net(nn.Module):
    def __init__(self, section_depth):
        super().__init__()
        #! define network architecture
        self.section_depth = section_depth
        self.num_filters = round(16/np.sqrt(self.section_depth))
        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.avgpool = nn.AvgPool2d(kernel_size=8)
        self.block1 = nn.ModuleList()
        self.block2 = nn.ModuleList()
        self.block3 = nn.ModuleList()
        self.batchnorm1 = nn.BatchNorm2d(self.num_filters)
        self.batchnorm2 = nn.BatchNorm2d(2*self.num_filters)
        self.batchnorm3 = nn.BatchNorm2d(4*self.num_filters)
        for i in range(self.section_depth):
            channels1 = 3 if i==0 else self.num_filters
            channels2 = self.num_filters if i == 0 else 2*self.num_filters
            channels3 = 2*self.num_filters if i == 0 else 4*self.num_filters
            self.block1.append(nn.Conv2d(in_channels=channels1, out_channels=self.num_filters, kernel_size=3, padding='same'))
            self.block2.append(nn.Conv2d(in_channels=channels2, out_channels=2*self.num_filters, kernel_size=3, padding='same'))
            self.block3.append(nn.Conv2d(in_channels=channels3, out_channels=4*self.num_filters, kernel_size=3, padding='same'))
        self.fc1 = nn.Linear(4*self.num_filters, 10)  # ? number of outputs
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        for i in self.block1:
            x = F.relu(self.batchnorm1(i(x)))
        x = self.maxpool(x)
        for i in self.block2:
            x = F.relu(self.batchnorm2(i(x)))
        x = self.maxpool(x)
        for i in self.block3:
            x = F.relu(self.batchnorm3(i(x)))
        x = self.avgpool(x)
        x = torch.flatten(x, 1) # flatten all dimensions except batch
        x = self.fc1(x)
        x = self.softmax(x)
        return x

Any help would be appreciated.

Answer 1

Ok, I figured it out. Turns out that batch normalization layer has learnable parameters, so instead of using the same layer for the whole convblock I had to make a separate batchnorm layer for every convolutional layer.