How to add layers to a pretrained model in PyTorch?

Question

I want to add layer normalization function just after AdaptiveAvgPool2d layer and L2 normalization after fc layer. i wanted my fc layer output to be 200 so tried not to include fc layer instead of it make new fc layer, but it did not remove fc layers comes with pretrained model, i am using googlenet.

my code:

class GoogleNet(nn.Module):
    def __init__(self):
        super(GoogleNet,self).__init__()
        self.model = googlenet_pytorch.GoogLeNet.from_pretrained('googlenet')
        
        self.fc = nn.Linear(1024,200, bias=False),
        
    def forward(self, x):
        batch_size ,_,_,_ =x.shape
        x = self.model.extract_features(x)
        x = self.model._avg_pooling(x)
        x = F.layer_norm(x,x.size[1:],elementwise_affine=False)
        x = self.fc(x)
        x = F.normalize(x, p=2, dim=1)
        return x

Output I am getting:

      .....  
      .....  
      .....

(aux1): InceptionAux(
      (conv): BasicConv2d(
        (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn): BatchNorm2d(128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
      )
      (fc1): Linear(in_features=2048, out_features=1024, bias=True)
      (fc2): Linear(in_features=1024, out_features=1000, bias=True)
    )
    (aux2): InceptionAux(
      (conv): BasicConv2d(
        (conv): Conv2d(528, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn): BatchNorm2d(128, eps=0.001, momentum=0.1, affine=True, track_running_stats=True)
      )
      (fc1): Linear(in_features=2048, out_features=1024, bias=True)
      (fc2): Linear(in_features=1024, out_features=1000, bias=True)
    )
    (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
    (dropout): Dropout(p=0.2, inplace=False)
    (fc): Linear(in_features=1024, out_features=1000, bias=True)
  )
)

Output I want:

  ......
  ......
  ......

  (aux1): None
  (aux2): None
  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))
  ** layer normalization here**
  (dropout): Dropout(p=0.2, inplace=False)
  (fc): Linear(in_features=1024, out_features=200, bias=False)
  **L2 normalization here**
)

If someone needs to know the solution code for this problem, with the help of iacob's answer I solved it, I added it as an answer.

Answer 1

When printing a model, PyTorch will print every layer up until it encounters an error in the forward, irrespective of whether the model will run (on appropriately formatted input data).

You have a number of issues in your code after loading the backbone GoogLeNet model, and hence all layers you have added after this fail to display when printed. The immediately obvious issues are:

• you must remove the , after self.fc = nn.Linear(1024,200, bias=False), or it will be interpreted as a tuple
• GoogLeNet has no attribute _avg_pooling
• x.size is not subscriptable. Either use a function call x.size()[1:] or use x.shape[1:]
• F.layer_norm() has no argument elementwise_affine

You need to correct these in order to get the model to run, like so.

Answer 2

I think you can just remove the last layers and then add the layers you want.

So in your case:

class GoogleNet(nn.Module):
    def __init__(self):
        super(GoogleNet,self).__init__()

        # load the original google net
        self.model = googlenet_pytorch.GoogLeNet.from_pretrained('googlenet')
        
        # remove the last two layers (fc and dropout)
        self.model = nn.Sequential(*list(self.model.children())[:-2])
        
        # add your dropout layer
        self.dropout = nn.Dropout(0.2, inplace=False)

        # add your dense layer
        self.fc = nn.Linear(1024, 200, bias=False)
        
    def forward(self, x):
        batch_size, _, _, _ = x.shape

        # I dont know what this does but I'll leave it here
        x = self.model.extract_features(x)
        # but I would add x = self.model(x) instead

        # insert your layer normalization
        x = F.layer_norm(x,x.size[1:],elementwise_affine=False)
        
        # put droput layer back on
        x = self.dropout(x)
        
        x = self.fc(x)

        x = F.normalize(x, p=2, dim=1)

        return x

But note that when you print the model summary it will not show you the normalization layers. It doesnt print the stuff you add using F. in forward(), only the initialiazions you made in __init__().

Answer 3

After doing corrections in code according to @iacob i solved the problem.

correct code:

class GoogleNet(nn.Module):
    def __init__(self):
        super(GoogleNet,self).__init__()
        self.model = googlenet_pytorch.GoogLeNet.from_pretrained('googlenet')
        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
        self.layernorm = nn.LayerNorm(1024,elementwise_affine=True)
        self._fc = nn.Linear(1024,200, bias=False)
    def forward(self, x):
        batch_size ,_,_,_ =x.shape
        x = self.model.extract_features(x)
        x = self.model.avgpool(x)
        x = x.view(-1, 1024)
        x = self.layernorm(x)
        x = self._fc(x)
        x = F.normalize(x, p=2, dim=1)
        return x

After this i make an instance of GoogleNet() and converts the layers that i do not need to Identity().

dep = GoogleNet()
dep.model.aux1 = Identity()
dep.model.aux2 = Identity()
dep.model.avgpool = Identity()
dep.model.dropout = Identity()
dep.model.fc = Identity()