I am currently working on implementation of this paper Fully Convolutional Networks for Semantic Segmentation 4. I am first trying to build a FCN-32 architecture. I am using VGG16 pretrained model for training PASCAL VOC 2012 dataset  4. I am having a doubt regarding how should I use pretrained weights of VGG16 in my custom class implementation’s feature’s section. I have created a code snippt which demonstrates the above scenario. Can anyone tell me if this implementation of using pretrained weights in custom model correct ?

#loading the pretrained VGG16 model

model1 = models.vgg16(pretrained=True)

 

#Freezing the layers except the fc layers

for param in model1.features.paramters():

    param.requires_grad = False

   

#Creating FCN custom module:

class FCN(nn.Module):

    def __init__(self):

        super(FCN,self).__init__()

        self.features = nn.Sequential(*list(model1.features.children()))

        self.classifier = nn.Sequential(nn.Conv2d(512,4096,7),

                            nn.Dropout(),

                            nn.Conv2d(4096,21,1),

                            nn.Dropout(),

                            nn.ConvTranspose2d(21,21,224,stride=32)

                        )

    def forward(self,x):

        x = self.features(x)

        x = self.classifier(x)

        return x

       

model2 = FCN()

 

#Again freezing the feature's layer of model2

for params in model2.features.parameters():

    params.requires_grad = False

   

#For confirming that all the pretrained weights from vgg16 are transfered to FCN custom model. must return true for confirmation

print(list(model2.features.parameters()) == list(model1.features.parameters()))