tests

Author: Borda

pytorch_lightning/callbacks/model_checkpoint.py

@@ -335,6 +335,7 @@ def _format_checkpoint_name(
         epoch: int,
         step: int,
         metrics: Dict[str, Any],
+        prefix: str = "",
     ) -> str:
         if not filename:
             # filename is not set, use default name
@@ -351,6 +352,9 @@ def _format_checkpoint_name(
                     metrics[name] = 0
             filename = filename.format(**metrics)
 
+        if prefix:
+            filename = cls.CHECKPOINT_JOIN_CHAR.join([prefix, filename])
+
         return filename
 
     def format_checkpoint_name(self, epoch: int, step: int, metrics: Dict[str, Any], ver: Optional[int] = None) -> str:

tests/checkpointing/test_model_checkpoint.py

@@ -294,9 +294,9 @@ def test_model_checkpoint_format_checkpoint_name(tmpdir):
     assert ckpt_name == str(Path('.').resolve() / 'epoch=3-step=4.ckpt')
 
     # with version
-    ckpt = ModelCheckpoint(monitor='early_stop_on', dirpath=tmpdir, filename='name', prefix='test')
+    ckpt = ModelCheckpoint(monitor='early_stop_on', dirpath=tmpdir, filename='name')
     ckpt_name = ckpt.format_checkpoint_name(3, 2, {}, ver=3)
-    assert ckpt_name == tmpdir / 'test-name-v3.ckpt'
+    assert ckpt_name == tmpdir / 'name-v3.ckpt'
 
     # using slashes
     ckpt = ModelCheckpoint(monitor='early_stop_on', dirpath=None, filename='{epoch}_{val/loss:.5f}')

tests/trainer/test_trainer.py

@@ -421,34 +421,17 @@ def test_dp_output_reduce():
 
 
 @pytest.mark.parametrize(
-    ["save_top_k", "save_last", "file_prefix", "expected_files"],
+    ["save_top_k", "save_last", "expected_files"],
     [
-        pytest.param(
-            -1,
-            False,
-            "",
-            {"epoch=4.ckpt", "epoch=3.ckpt", "epoch=2.ckpt", "epoch=1.ckpt", "epoch=0.ckpt"},
-            id="CASE K=-1  (all)",
-        ),
-        pytest.param(1, False, "test_prefix", {"test_prefix-epoch=4.ckpt"}, id="CASE K=1 (2.5, epoch 4)"),
-        pytest.param(2, False, "", {"epoch=4.ckpt", "epoch=2.ckpt"}, id="CASE K=2 (2.5 epoch 4, 2.8 epoch 2)"),
-        pytest.param(
-            4,
-            False,
-            "",
-            {"epoch=1.ckpt", "epoch=4.ckpt", "epoch=3.ckpt", "epoch=2.ckpt"},
-            id="CASE K=4 (save all 4 base)",
-        ),
-        pytest.param(
-            3,
-            False,
-            "", {"epoch=2.ckpt", "epoch=3.ckpt", "epoch=4.ckpt"},
-            id="CASE K=3 (save the 2nd, 3rd, 4th model)"
-        ),
-        pytest.param(1, True, "", {"epoch=4.ckpt", "last.ckpt"}, id="CASE K=1 (save the 4th model and the last model)"),
+        pytest.param(-1, False, [f"epoch={i}.ckpt" for i in range(5)], id="CASE K=-1  (all)"),
+        pytest.param(1, False, {"epoch=4.ckpt"}, id="CASE K=1 (2.5, epoch 4)"),
+        pytest.param(2, False, [f"epoch={i}.ckpt" for i in (2, 4)], id="CASE K=2 (2.5 epoch 4, 2.8 epoch 2)"),
+        pytest.param(4, False, [f"epoch={i}.ckpt" for i in range(1, 5)], id="CASE K=4 (save all 4 base)"),
+        pytest.param(3, False, [f"epoch={i}.ckpt" for i in range(2, 5)], id="CASE K=3 (save the 2nd, 3rd, 4th model)"),
+        pytest.param(1, True, {"epoch=4.ckpt", "last.ckpt"}, id="CASE K=1 (save the 4th model and the last model)"),
     ],
 )
-def test_model_checkpoint_options(tmpdir, save_top_k, save_last, file_prefix, expected_files):
+def test_model_checkpoint_options(tmpdir, save_top_k, save_last, expected_files):
     """Test ModelCheckpoint options."""
 
     def mock_save_function(filepath, *args):
@@ -463,7 +446,6 @@ def mock_save_function(filepath, *args):
         monitor='checkpoint_on',
         save_top_k=save_top_k,
         save_last=save_last,
-        prefix=file_prefix,
         verbose=1
     )
     checkpoint_callback.save_function = mock_save_function