[FIXED] Modifying the forward function in vgg model

Issue

I need to modify the existing forward method in VGG16 so that it can pass through the two classifiers and return the value

I tried creating the custom forward method manually and overriding the existing method but I get the following error

vgg.forward = forward

forward() missing 1 required positional argument: 'x'

My custom forward function

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

I have modified the default vgg16_bn with one additional classifier as

vgg = models.vgg16_bn()
final_in_features = vgg.classifier[6].in_features
mod_classifier = list(vgg.classifier.children())[:-1]
mod_classifier.extend([nn.Linear(final_in_features, 10)])
vgg.add_module('classifier_2',vgg.classifier)

My model looks like this after the addition of above classifier

(classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace)
    (2): Dropout(p=0.5)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace)
    (5): Dropout(p=0.5)
    (6): Linear(in_features=4096, out_features=10, bias=True)
  )
  (classifier_2): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace)
    (2): Dropout(p=0.5)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace)
    (5): Dropout(p=0.5)
    (6): Linear(in_features=4096, out_features=10, bias=True)
  )

My convolutional layers results are supposed to be passed through two separate FFN layers. So how do i modify my forward pass

Solution

I think the best way to achieve what you want is to create a new model extending the nn.Module. I'd do something like:

from torchvision import models
from torch import nn

class MyVgg (nn.Module):

    def __init__(self):
        super(Net, self).__init__()

        vgg = models.vgg16_bn(pretrained=True)

        # Here you get the bottleneck/feature extractor
        self.vgg_feature_extractor = nn.Sequential(*list(vgg.children())[:-1])

        # Now you can include your classifiers
        self.classifier1 = nn.Sequential(layers1)
        self.classifier2 = nn.Sequential(layers2)

    # Set your own forward pass
    def forward(self, img, extra_info=None):

        x = self.vgg_convs(img)
        x = x.view(x.size(0), -1)
        x1 = self.classifier1(x)
        x2 = self.classifier2(x)

        return x1, x2

Answered By - André Pacheco