Define Network (VGG) as one Sequential layer

vision
1

Hello,

I’m trying to transfer the torchvision.models.vgg11 into one Sequential layer. You can take a look at the PyTorch source-code here.
As you can see the forward pass is defined as:

def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

So I thought it could work like this:

pt_vgg11 = vgg11(pretrained=False)

class FlattenModule(nn.Module):
    def forward(self, x):
        x = torch.flatten(x, 1)
        return x

def VGG11_Seq():
    features = torch.nn.Sequential(
        pt_vgg16.features,
        pt_vgg16.avgpool,
        FlattenModule(),
        pt_vgg16.classifier
    )
    return features

model = VGG11_Seq()

(I also tried to print out the pt_vgg11 and just take everything that’s part of the network and put it inside of a Sequential, but that won’t work either …)

but when I try to load weights from a trained torchvision.models.vgg11
it won’t work. And also training the model = VGG11_Seq() won’t work. (working with CIFAR10)

Can anyone see the error, or can provide any tips?

2

To define the vgg model as one sequential layer you could simply do -

vgg = models.vgg16(true/false)
vgg_features = nn.Sequential(*list(vgg.children())[:-1]) #removed the Last linear layer

And in the forward method -
out = vgg_features(input)

1 Like
3

Thank you for your answer! Why do you remove the last linear layer?

4

You are extracting features from the pre-trained VGG and training your own classification layer, I am right to assume that? Because you are extracting features then flattening them. If you do not remove the last layer of the VGG, your output would be of the shape-[batch_size, num_classes], already giving you the class probabilities.
Though you can always add extra linear layers on top of the VGG network.

1 Like