Picking the best model from continuing runs

[youngjae142]

Hi,
I just realized that when a training continues from the previously truncated run, it picks the best model only from the current run?
In the figure below, the training was truncated at the 11th epoch, and continued.
The continued run started from the 12th epoch, and started to re-pick every epoch as the best model.
(which means it cannot remember the best model from the previous run.)

Is there a way to enforce picking the best model among all runs after initialization?
My NequIP version is 0.6.1, FYI.

[cw-tan]

Hi @youngjae142 ,

Sorry but v0.6.1 is very old and will not be maintained. Otherwise, happy to answer questions about more recent versions.

[youngjae142]

Hi @cw-tan,
Thanks for the reply!
I am not expecting a maintenance for this version for sure! 😅
I wonder this ``best model picking'' behavior is now different in recent versions? (like picking the best one from all runs after initialization?

[cw-tan]

See https://nequip.readthedocs.io/en/latest/guide/getting-started/workflow.html#saving-and-restarting

For continuing an interrupted run, you have to provide the path to the specific checkpoint file you want to continue from. It is user-controlled in that sense.

[youngjae142]

I know I can start from a specific checkpoint.

I am talking about the fact that the "restarting" is not recovering the "best model" information from the last run.

[cw-tan]

Ah I see, sorry. I'm not sure actually. Looking at the load_state_dict hook of the ModelCheckpoint callback from lightning (https://lightning.ai/docs/pytorch/stable/_modules/lightning/pytorch/callbacks/model_checkpoint.html#ModelCheckpoint):

[[docs]](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#lightning.pytorch.callbacks.ModelCheckpoint.load_state_dict)    @override
    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        dirpath_from_ckpt = state_dict.get("dirpath", self.dirpath)

        if self.dirpath == dirpath_from_ckpt:
            self.best_model_score = state_dict["best_model_score"]
            self.kth_best_model_path = state_dict.get("kth_best_model_path", self.kth_best_model_path)
            self.kth_value = state_dict.get("kth_value", self.kth_value)
            self.best_k_models = state_dict.get("best_k_models", self.best_k_models)
            self.last_model_path = state_dict.get("last_model_path", self.last_model_path)
        else:
            warnings.warn(
                f"The dirpath has changed from {dirpath_from_ckpt!r} to {self.dirpath!r},"
                " therefore `best_model_score`, `kth_best_model_path`, `kth_value`, `last_model_path` and"
                " `best_k_models` won't be reloaded. Only `best_model_path` will be reloaded."
            )

        self.best_model_path = state_dict["best_model_path"]

It always does self.best_model_path = state_dict["best_model_path"], which may imply that it retains memory of the best model from the previous run. But one can do some tests to verify this.

[youngjae142]

Oh, the "best model" picking function is from PyTorch? I assumed it is a function of nequip.
I'll look into it. Thanks!