Add best_k_metrics parameter to the ModelCheckpoint
#20457
Conversation
gonzachiar
requested review from
lantiga,
Borda,
tchaton,
justusschock and
ethanwharris
as code owners
| @@ -241,9 +249,10 @@ def __init__( | |||
| self._last_global_step_saved = 0 # no need to save when no steps were taken | |||
| self._last_time_checked: Optional[float] = None | |||
| self.current_score: Optional[Tensor] = None | |||
| self.best_k_models: Dict[str, Tensor] = {} | |||
| self.best_k_models: Dict[str, Dict[str, Tensor | Dict[str, Tensor]]] = {} | |||
There was a problem hiding this comment.
this may be easier to read
Dict[str, Dict[str, Tensor]] | Dict[str, Dict[str, Dict[str, Tensor]]]
but ultimately we'd be better off defining a type alias
more importantly, we need to avoid breaking backward compatibility here
so whatever code relies on best_k_models being Dict[str, Tensor] today needs to keep working
I suggest we just limit ourselves to track best_model_metrics and not mess with best_k_models, or use a separate private attribute
| @@ -523,7 +534,9 @@ def check_monitor_top_k(self, trainer: "pl.Trainer", current: Optional[Tensor] = | |||
| return True | |||
|
|
|||
| monitor_op = {"min": torch.lt, "max": torch.gt}[self.mode] | |||
| should_update_best_and_save = monitor_op(current, self.best_k_models[self.kth_best_model_path]) | |||
| should_update_best_and_save = monitor_op( | |||
| current, cast(Tensor, self.best_k_models[self.kth_best_model_path]["score"]) | |||
There was a problem hiding this comment.
this will stay as in the original if we avoid changing best_k_models
| @@ -706,6 +706,7 @@ def test_model_checkpoint_save_last_none_monitor(tmp_path, caplog): | |||
| assert checkpoint_callback.best_model_path == str(tmp_path / "epoch=1-step=20.ckpt") | |||
| assert checkpoint_callback.last_model_path == str(tmp_path / "last.ckpt") | |||
| assert checkpoint_callback.best_model_score is None | |||
| assert checkpoint_callback.best_model_metrics is None | |||
There was a problem hiding this comment.
we need to add tests that exercise the new code
|
@gonzachiar I'm wrapping up the last few PRs for the release, do you have time to push this through in the next couple of days? |
for more information, see https://pre-commit.ci
Hey, I will work on this during the weekend. Hopes that helps! |
|
This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 7 days if no further activity occurs. If you need further help see our docs: https://lightning.ai/docs/pytorch/latest/generated/CONTRIBUTING.html#pull-request or ask the assistance of a core contributor here or on Discord. Thank you for your contributions. |
for more information, see https://pre-commit.ci
Borda
changed the title
best_k_metrics parameter to the ModelCheckpoint
What does this PR do?
Adds a parameter to save all the metrics from the best model.
Fixes #20321
Before submitting
PR review
Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the review guidelines. In short, see the following bullet-list:
Reviewer checklist
📚 Documentation preview 📚: https://pytorch-lightning--20457.org.readthedocs.build/en/20457/