How to remove the fully connected layers of my pretrained VGG net

Hello,

this is my first post in that forum and I have the following problem/question.

I have trained a VGG11 net to do a binary classification and now I want to use the pretrained net in another way, too. I want to use the pretrained net without the fully connected layers for an image segmentation task.

cfg = {
    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}


class VGG(nn.Module):
    def __init__(self, vgg_name):
        super(VGG, self).__init__()
        self.features = self._make_layers(cfg[vgg_name])
        self.classifier = nn.Linear(512, 2)

    def forward(self, x):
        out = self.features(x)
        out = out.view(out.size(0), -1)
        out = self.classifier(out)
        return out

    def _make_layers(self, cfg):
        layers = []
        in_channels = 3
        for x in cfg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
                           nn.BatchNorm2d(x),
                           nn.ReLU(inplace=True)]
                in_channels = x
        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
        return nn.Sequential(*layers)


net = VGG('VGG11')

Thats how I initializes the net(Got that from GitHub). After training I save it with torch.save().

torch.save(net, 'VGG11.pt')

How can I load the ‘VGG11.pt’ without the fully connected layers now?

Im looking forward for your answers, best regards.

(Im sorry for my bad English)

First of all it is recommended not to use torch.save to save your model. Hav a look at this page for more details.

For your question: I’m afraid you have to remove them before saving or save the whole model and load the whole model (you could still remove the layers after loading).

Do you want to add other FC Layers instead?

1 Like

First of all, ty for your advice.

I would have no problem with removing them before saving if that would work.

I want to add a fully convolutional net on top of the feature extracting part of the pretrained vgg net. That I can do image segmentation. I already found something like this on GitHub. (here) In this case the author uses a pretrained vgg net with out the fc layers as base. I tried this implementation with my own dataset with out success. The results are not as expected.

I think the problem is that the pretrained vgg model does not fit that good for my specific case (Segmentation of PCBs). Therefore I have trained my own vgg net that can classificate different parts of a PCB. (Coil, Microchip…) And that pretrained vgg net I want to us now as base for the fully convolutional net and try if I can reach better results.

Then why don’t you just save net.features?

I’m not sure if you have to wrap it in an own module, but you could do something like this:

class VGGFeatures(nn.Module):
    def __init__(self, vgg_net: VGG):
        super().__init__()
        self.vgg_features = vgg_net.features
        # maybe add your segmentation part

    @classmethod
    def from_state_dict(cls, state_dict_path, vgg_type='VGG11'):
        _net = VGG(vgg_type)
        _current_instance = cls(_net)
        _current_instance.load_state_dict(torch.load(state_dict_path))
        return _current_instance
        

    def forward(x):
        features = self.vgg_features(x)
        # maybe forward through your segmentation part
        return features

net = VGG('VGG11')
vgg_features = VGGFeatures(net)

# save model
torch.save(vgg_features.state_dict(), SAVE_PATH)

# recover model
recovered = VGGFeatures.from_state_dict(SAVE_PATH, 'VGG11')

I saved my net now the way you told me to do it. You can see it below.

torch.save(net.state_dict(),'VGGNet11.pt')

And now I want to load the net in another python file. The Pytorch help page says that you have to load the net like this:

net = TheModelCLass(*args, **kwargs)
net.load_state_dict(torch.load('VGGNet11.pt'))

What do they mean with TheModelCLass?

TheModelClass would be VGGFeatures in your case. You have to create a model with randomly initialized weights and afterwords replace these random weights with your saved weights to keep it directory independent.

Have you tried this?

vgg11 = torchvision.models.vgg11()
vgg11.load_state_dict(torch.load('path/to/your/saved/vgg11/weights.pt'))

vgg11_feature_detector = vgg11.features
new_conv_layer = nn.Conv2d(num_in, num_out, kernel, stride)

my_net = nn.Sequential(*[vgg11_feature_detector, new_conv_layer])

Okay I could now load the model.

But for some reasons I get different accuracy with the loaded weights. I used the same test images in the same order. Im pretty confused now.

Thats my test method:

def test(test_data):
    """
            Testfunktion
    """
    correct = 0
    total = 0
    for data in test_data:
        images, labels = data
        labels = Variable(torch.Tensor(labels))
        images = Variable(images)
        outputs = net(images)

        predicted = outputs.data.max(1, keepdim=True)[1]
        labels = labels.data.max(1, keepdim=True)[1]

        # Bestimmung der Anzahl der Testdaten
        total += labels.size(0)

        # Aufsummieren der korrekten Vorhersagen
        correct += (predicted == labels).sum()

    z.write('Accuracy of the network on the %d test images after epoche %d: %.2f %%\n' % (total, epoch+1, 100 * correct / total))

And I call the method after every epoch of training here:

for epoch in range(epochen):
    """ 
        Trainingsprozess in Epochen
    """
    # Netz wird auf Trainieren gestellt
    net.train()

    running_loss = 0.0
    for i, data in enumerate(train_data, 0):

        # Daten und Label der Traingsdaten
        inputs, labels = data

        # Zusammenfassen der Eingabetensoren in Variablen. ( Wird benötigt für den rückwärtsgerichteten Weg)
        inputs, labels = Variable(inputs), Variable(torch.Tensor(labels))

        # Zurücksetzen der Gradienten.
        optimizer.zero_grad()

        # forward + backward + optimieren(Gewichte updaten)
        outputs = F.softmax(net(inputs),dim=1)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # Darstellung einer Statstik
        running_loss += loss.data

        # Der Fehler jedes 50. Batches wird ausgegeben.
        if i % 50 == 49:
            print('[Epoche: %d Batch: %d] loss: %.3f' %
                  (epoch + 1, i + 1, running_loss / 50))
            running_loss = 0.0

    # Netz wird auf Testen gestellt. (Gewichte sind fest)
    net.eval()

    # Aufrufen der Testfunktion
    test(test_data)

Did you save/load it including the FC layer or did you add another FC layer afterwards?

with the fc layers. I just want to test loading and saving the normal net for now.

To test the loaded net. I load the net and call the test method with the same test_data I tested the net before.

I tried a lot not but I can’t find any problem.

Thats the net class I use in the python file I train the model:

cfg = {
    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}


class VGG(nn.Module):
    def __init__(self, vgg_name):
        super(VGG, self).__init__()
        self.features = self._make_layers(cfg[vgg_name])
        self.classifier = nn.Linear(512, 2)

    def forward(self, x):
        out = self.features(x)
        out = out.view(out.size(0), -1)
        out = self.classifier(out)
        return out

    def _make_layers(self, cfg):
        layers = []
        in_channels = 3
        for x in cfg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
                           nn.BatchNorm2d(x),
                           nn.ReLU(inplace=True)]
                in_channels = x
        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
        return nn.Sequential(*layers)

# Initialisieren des Netzes
net = VGG('VGG11')

And I save it like that:

# trainiertes Netzspeichern
torch.save(net.state_dict(),'VGGNet11.pt')

In the python file I want to test the trained net I load it like that:

cfg = {
    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],
    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],
}


class VGG(nn.Module):
    def __init__(self, vgg_name):
        super(VGG, self).__init__()
        self.features = self._make_layers(cfg[vgg_name])
        self.classifier = nn.Linear(512, 2)

    def forward(self, x):
        out = self.features(x)
        out = out.view(out.size(0), -1)
        out = self.classifier(out)
        return out

    def _make_layers(self, cfg):
        layers = []
        in_channels = 3
        for x in cfg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
            else:
                layers += [nn.Conv2d(in_channels, x, kernel_size=3, padding=1),
                           nn.BatchNorm2d(x),
                           nn.ReLU(inplace=True)]
                in_channels = x
        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]
        return nn.Sequential(*layers)

# Initialisieren des Netzes
net = VGG('VGG11')

net.load_state_dict(torch.load('VGGNet11.pt'))

Both python files are located in the same folder.

I tested both nets(the loaded and the pre-loaded after training) on one images now and both nets compute a different output. Thats means the mistake has to happen while loading or saving. Im pretty confused now.

Thats the method I tested both nets with:

def wtf():
    net.eval()
    img = Image.open('file path....')
    img_tensor = Variable(transform(img)).unsqueeze_(0)
    outputs = net(img_tensor)
    print(outputs)

I found out something new. I tried something in my training python file. I trained the net and called the method wtf. Then I saved the net and loaded it straight again and called the wtf method with the loaded net again. Now I get the same outputs from both nets. Now im totally confused but that means that the problem is not to save the net. The problem is to load the net in the new python file.

Thats the code I run after training:

def wtf():
    net.eval()
    img = Image.open('path...')
    img_tensor = Variable(transform(img)).unsqueeze_(0)
    outputs = net(img_tensor)
   print('output net: {}'.format(outputs))
def wtf2():
    loaded_net.eval()
    img = Image.open('path...')
    img_tensor = Variable(transform(img)).unsqueeze_(0)
    outputs = loaded_net(img_tensor)
    print('output loaded_net: {}'.format(outputs))

wtf()
torch.save(net.state_dict(),'VGGNet11.pt')
loaded_net = VGG('VGG11')
loaded_net.load_state_dict(torch.load('VGGNet11.pt'))
wtf2()

output net: Variable containing:
-1.3904  1.3422
[torch.FloatTensor of size 1x2]

output loaded_net: Variable containing:
-1.3904  1.3422
[torch.FloatTensor of size 1x2]

But the output seems to be identical?