Pytorch: Lower the parameters in U-net model
can anyone give me some tips on how i would be able to lower the amount of parameters in the following U-net implementation. I'm having trouble with over-fitting on my training data and i would like to lower the parameters in order to see if it improves the validation data accuracy.
Layers:
First2D
layers = [
        nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(middle_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    ]

Encoder2D
layers = [
        nn.MaxPool2d(kernel_size=downsample_kernel),
        nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(middle_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    ]

Center2D
layers = [
        nn.MaxPool2d(kernel_size=2),
        nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(middle_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(out_channels, deconv_channels, kernel_size=2, stride=2)
    ]

Decoder2D
layers = [
        nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(middle_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        nn.ConvTranspose2d(out_channels, deconv_channels, kernel_size=2, stride=2)
    ]

Last2D
layers = [
        nn.Conv2d(in_channels, middle_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(middle_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(middle_channels, middle_channels, kernel_size=3, padding=1),
        nn.BatchNorm2d(middle_channels),
        nn.ReLU(inplace=True),
        nn.Conv2d(middle_channels, out_channels, kernel_size=1),
        nn.Softmax(dim=1)
    ]


One way to decrease the number of parameters is to decrease the number of channels in the convolution. You wouldn't be able to change the number of model input and output channels, because they depend on the data, but you can change the number of intermediate channels.
Remember that the output of one layer is the input to the next layer, so keep the number of output channels in the first layer the same as the number of input channels in the second layer, for every pair of layers. Example would be
layers = [
    nn.Conv2d(in_channels, middle_channels//2, kernel_size=3, padding=1),
    nn.BatchNorm2d(middle_channels//2),
    nn.ReLU(inplace=True),
    nn.Conv2d(middle_channels//2, out_channels, kernel_size=3, padding=1),
    nn.BatchNorm2d(out_channels),
    nn.ReLU(inplace=True)
]

Now, coming to the original question of overfitting, first you might want to try to use other things first, before reducing model size. Some things include data augmentations and dropout.
