trypag/pytorch-unet-segnet

Incorrect number of convolutional layers in SegNet model

Opened this issue · 0 comments

glhr commented

In the encoder and decoder blocks, there can never be more than 2 convolutional layers, even if the n_blocks argument is > 2, since there is never more than a single layer added:

        layers = [nn.Conv2d(n_in_feat, n_out_feat, 3, 1, 1),
                  nn.BatchNorm2d(n_out_feat),
                  nn.ReLU(inplace=True)]

        if n_blocks > 1:
            layers += [nn.Conv2d(n_out_feat, n_out_feat, 3, 1, 1),
                       nn.BatchNorm2d(n_out_feat),
                       nn.ReLU(inplace=True)]
            if n_blocks == 3:
                layers += [nn.Dropout(drop_rate)]

So this means that encoder_n_layers = (2, 2, 3, 3, 3) does not actually ever produce 3-layer blocks. It only adds drop-out. Same issue for the decoder. So this model doesn't match the VGG/SegNet model.
Instead, it should be something like:

        layers = []
        for layer in range(n_blocks):
            layers += [nn.Conv2d(n_in_feat if layer == 0 else n_out_feat, n_out_feat, 3, 1, 1),
                       nn.BatchNorm2d(n_out_feat),
                       nn.ReLU(inplace=True)]
        if n_blocks > 2:
                layers += [nn.Dropout(drop_rate)]