Loading old checkpoints with new hparam keys

[codezakh]

Define a LightningModule with an __init__ like this....

class LitModel(pl.LightningModule):
    def __init__(self, cfg: DictConfig):
        self.backbone = nn.Linear(2,1)
        self.save_hyperparameters(cfg)

where cfg is say,

foo: bar

and then save the model with trainer.save_checkpoint() to model.ckpt. Then at a later time, we make a small change to LitModel and gate it behind a new key, so it is now:

class LitModel(pl.LightningModule):
    def __init__(self, cfg: DictConfig):
        self.backbone = nn.Linear(2,1)
        self.save_hyperparameters(cfg)
        # some new behavior
        if self.hparams.new_key:
            self.backbone.bias = self.backbone.bias * 0.5

and we attempt to load our saved model.ckpt with LitModel.load_from_checkpoint('model.ckpt', new_key=True) to test how the new behavior affects the performance of the old saved model.

The current behavior is to throw omegaconf.errors.ConfigKeyError: Key 'new_key' is not in struct because of this_line, which is part of the following block:

# for past checkpoint need to add the new key
if cls.CHECKPOINT_HYPER_PARAMS_KEY not in checkpoint:
    checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY] = {}
# override the hparams with values that were passed in
checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)

TLDR: Once a model has been saved with a DictConfig as its hyperparameter storage, it is seems impossible (please correct me if I'm wrong) to load it with future code that adds more hparams, even if they are compatible with the architecture.
What is the recommended way to handle this use case, because it is natural over the course of R&D to add new flags and want to load old weights with them. My current solution is to override load_from_checkpoint and add OmegaConf.set_struct(..., False) before checkpoint[cls.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs) or use the omegaconf.open_dict(...) context manager before the .update() call.
It feels as if lightning should do this, because the hparams object on an instantiated LightningModule does not seem to have the struct=True flag set on it.

[dasayan05]

@codezakh, The .load_from_checkpoint(..) basically discards new hparams to be injected into the constructor.

I initiated a feature request #8631 (@tchaton agreed on the idea) few days back which is a slightly different use case but requires pretty much similar fix. Not sure if it has any drawbacks though.