Problem loading pretrained model, fine tune model on custom dataset

vision
1

Step 1 : I am loading pretrained xception model on ImageNet and fine-tuning for binary classification on my dataset.

from pretrainedmodels.models.xception import Xception
self.xception_model = Xception(num_classes=1000)
.
.
.
.
checkpoint = torch.load('/home/shivangi/Desktop/Projects/pretrained_models/checkpoints/xception-43020ad28.pth')
self.xception_model.load_state_dict(checkpoint)
self.xception_model.last_linear = nn.Linear(2048, 2)

Step 2 : Now after training is complete I save this trained model. Everything works fine.

torch.save(self.xception_model.state_dict(), model_path  + self.model_name + '.pt)

Step 3 : Now I want use this model in step 2 on a different dataset

self.xception_model = Xception(num_classes=2)
checkpoint = torch.load(model_path  + self.model_name + '.pt')
self.xception_model.load_state_dict(checkpoint)

Then I get the error : AttributeError: ‘Xception’ object has no attribute ‘last_linear’

2

It looks like the actual factory to create the model is defiend here, i.e. as pretrainedmodels.models.xception(), while you created an instance of xception.Xception manually.

Since self.last_linear is only redefined in the factory (line of code), you could try to set the last_linear again to your custom linear layer:

self.xception_model = Xception(num_classes=2)
checkpoint = torch.load(model_path  + self.model_name + '.pt')
self.xception_model.last_linear = nn.Linear(2048, 2)
self.xception_model.load_state_dict(checkpoint)