Add ModelCheckpoint.to_yaml method

pytorch_lightning/callbacks/model_checkpoint.py

@@ -22,8 +22,10 @@
 
 import os
 import re
+import yaml
 from copy import deepcopy
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Union
+from pathlib import Path
 
 import numpy as np
 import torch
@@ -531,3 +533,13 @@ def _do_check_save(
             if cur_path != filepath:
                 self._del_model(cur_path)
 
+    def to_yaml(self, filepath: Optional[Union[str, Path]] = None):
+        """
+        Saves the `best_k_models` dict containing the checkpoint
+        paths with the corresponding scores to a YAML file.
+        """
+        best_k = {k: v.item() for k, v in self.best_k_models.items()}
+        if filepath is None:
+            filepath = os.path.join(self.dirpath, "best_k_models.yaml")
+        with open(filepath, "w") as fp:
+            yaml.dump(best_k, fp)

tests/callbacks/test_model_checkpoint.py

@@ -1,4 +1,5 @@
 import os
+import yaml
 import pickle
 import platform
 import re
@@ -36,6 +37,24 @@ def test_model_checkpoint_with_non_string_input(tmpdir, save_top_k):
     )
 
 
+@pytest.mark.parametrize('save_top_k', [-1, 0, 1, 2])
+def test_model_checkpoint_to_yaml(tmpdir, save_top_k):
+    """ Test that None in checkpoint callback is valid and that chkp_path is set correctly """
+    tutils.reset_seed()
+    model = EvalModelTemplate()
+
+    checkpoint = ModelCheckpoint(filepath=tmpdir, monitor='val_loss', save_top_k=save_top_k)
+
+    trainer = Trainer(default_root_dir=tmpdir, checkpoint_callback=checkpoint, overfit_batches=0.20, max_epochs=2)
+    trainer.fit(model)
+
+    path_yaml = os.path.join(tmpdir, 'best_k_models.yaml')
+    checkpoint.to_yaml(path_yaml)
+    d = yaml.full_load(open(path_yaml, 'r'))
+    best_k = {k: v.item() for k, v in checkpoint.best_k_models.items()}
+    assert d == best_k
+
+
 @pytest.mark.parametrize(
     "logger_version,expected",
     [(None, "version_0"), (1, "version_1"), ("awesome", "awesome")],