I am currently working on implementation of this paper Fully Convolutional Networks for Semantic Segmentation. I am first trying to build a FCN-32 architecture. I am using VGG16 pretrained model for training PASCAL VOC 2012 dataset . 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()))