Finetuning pre-trained model(add new loss function)

vision
1

Currently, my last bit of my model looks like this:

  (conv_last): Sequential(
    (0): Conv2d(2688, 5376, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (1): BatchNorm2d(5376, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace)
  )
  (globalpool): Sequential(
    (0): AvgPool2d(kernel_size=7, stride=7, padding=0)
  )
  (classifier): Sequential(
    (0): Linear(in_features=5376, out_features=2300, bias=True)
  )

with forward function:

    def forward(self, x):
        x = self.conv1(x)
        x = self.features(x)
        x = self.conv_last(x)
        x = self.globalpool(x)
        x = x.view(-1, self.stage_out_channels[-1])
        closs_output = self.linear_closs(x)
        x = self.classifier(x)
        return x

I wish to add a centre loss function, where the end result is like:


  (conv_last): Sequential(
    (0): Conv2d(2688, 5376, kernel_size=(1, 1), stride=(1, 1), bias=False)
    (1): BatchNorm2d(5376, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU(inplace)
  )
  (globalpool): Sequential(
    (0): AvgPool2d(kernel_size=7, stride=7, padding=0)
  )
  (classifier): Sequential(
    (0): Linear(in_features=5376, out_features=2300, bias=True)
  )
  (linear_closs): Linear(in_features=5376, out_features=1280, bias=False)
)

with new forward function:

    def forward(self, x):
        x = self.conv1(x)
        x = self.features(x)
        x = self.conv_last(x)
        x = self.globalpool(x)
        x = x.view(-1, self.stage_out_channels[-1])
        closs_output = self.linear_closs(x)
        x = self.classifier(x)
        return closs_output, x
#         return x

I used torch.load(model). Is there a way to add these new lines into the model? The only changes are the last layer of the model, and 2 new lines in the forward function to specify the return for the new loss function.

Thanks!

2

Hey! The model that you are loading is a key value dict with the key being the name of the layer and value the weights. You can add things to this dict, such as your new Linear layer. And yes, you can change the forward function to your liking :slight_smile:

There is some info on how to do the weights thingy in here

3

Thanks for the quick reply! Unfortunately I saved the model but not the state_dict, can I still make the changes?

4

You can save at a new path so you get to keep everything. You should be able to run your script multiple times until all the bugs are ironed out. So yes, you can still make the changes :confused:

5

I was able to add the last layer using:

network.add_module(module=nn.Linear(list(network.children())[-1][-1].in_features, 1024, bias = False), name='linear_closs')

However, I do not know how to change the forward function as described above. Went through the post but it did not cover that aspect. Any ideas? Thanks!

Note: the change to forward function as below:

  • Add additional line closs_output = self.linear_closs(x)
  • Change output of forward from return x to return closs_output, x
6

Hi,
I think you could not use network.add_module.
For an easy approach, you could just use torch.load(model) to load a network backbone and define your last linear layer in the new network. Additionally, the in_features=5376 of linear_closs layer you added should match out_features=2300 of the last layer in pre-trained model.
Here is a little snippet:

new_model = torch.load('model.pth')
input = torch.randn(1,3,3,3)
out = new_model(input)

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

        self.pre_trained_model = new_model
        self.last_linear = nn.Linear(27, 2)

    def forward(self, input):

        x = self.pre_trained_model(input)
        x = out.view(1,-1)
        closs_out = self.last_linear(x)

        return closs_out,x

model = New_network()
closs_out, out = model(input)
1 Like