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

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()