Using trainer.fit(model, datamodule, ckpt_path=path) on compressed model

[shenoynikhil98]

Task: Model Compression using NNI.

Approach: Loading a PyTorch Lightning trained model from a model checkpoint using .load_checkpoint().

Problem: After removing weights, my model class has reduced weights and has to be fine-tuned. As this is a compressed version of a trained model, I want to continue training with the optimizer state dict present in the checkpoint. If I try to fine-tune using

trainer = pl.Trainer(
     gpus=self.gpus,
     max_epochs=self.max_epochs,
     callbacks=[checkpoint_callback],
)
trainer.fit(self.model, datamodule, ckpt_path=ckpt_PATH)

This gives an error as the model structure has changed (Conv2d filters have been reduced). Is there an approach that can be used that is easier?

[rohitgr7]

this is not possible with fit(..., ckpt_path=...) since this is partial resume and will load the model weights too. For you use-case you can do this maybe:

class OptimizerReload(Callback):
    def __init__(self, ckpt_path):
        self.ckpt_path = ckpt_path

    def on_train_start(self, trainer, pl_module):
        ckpt = torch.load(self.ckpt_path)
        trainer.strategy.load_optimizer_state_dict(ckpt)

and, pass it to Trainer

cb = OptimizerReload(ckpt_path)
trainer = Trainer(..., callbacks=[cb])
trainer.fit(model)

[shenoynikhil98]

for <v1.6 try:
trainer.training_type_plugin.load_optimizer_state_dict(ckpt)
Again suggested by @rohitgr7 !