Is it possible to change the output dimension of a checkpoint model?

[shrijan00]

Hi,
I know this topic has been covered in the Pytorch forum but just wanted to know the best possible way to do it in Lightning.
So, I have a checkpoint for a model which gives 3 output but now I want to change it to 4. Is that easy to do in Lightning?

res_model = ResNet34(num_inputs_channels, numOutputs)
path = 'path to checkpoint'
trainer = pl.Trainer(max_epochs=config.nepochs, gpus = '0' , logger=False, callbacks=[checkpoint_callback], resume_from_checkpoint = path)
trainer.fit(model=res_model,datamodule=dataset)

How can I incorporate this change?

[tchaton]

Dear @shrijan00,

You could pass strict=False while loading the weights.

class MyModel(LightningModule):

    def __init__(self, ...):

        self.save_hyperparameters()

MyModel.load_from_checkpoint(chkpt_path, strict=False)

[shrijan00]

Dear @tchaton
Thanks for your reply. What I do now, take the trained model that had an output of 3, load the model and then change the final layer to output 4.

model_path = '20210713-ResNet34classes.pt'
if use_wandb == 'disabled':
    trainer = pl.Trainer(max_epochs=config.nepochs, gpus = '0' , logger=wandb_logger, callbacks=[checkpoint_callback], accumulate_grad_batches=3, resume_from_checkpoint=path)
else:
    trainer = pl.Trainer(max_epochs=config.nepochs, gpus = '0' , logger=wandb_logger, callbacks=[checkpoint_callback], accumulate_grad_batches=3 )
    
res_model.load_state_dict(torch.load(model_path))     
res_model.model.fc = nn.Sequential(nn.Dropout(0.5), nn.Linear(512, numOutputs))                                                                                         
trainer.fit(model=res_model,datamodule=dataset)

Here numPutputs is 4. Is this approach correct?

[tchaton]

Yes, this is the right approach.