Update save_for_pretrained to avod crashes when a matching checkpoint is not found and output details

flexthink · flexthink · commit 07155e9ec3e1 · 2022-06-26T15:46:00.000-04:00
diff --git a/speechbrain/pretrained/training.py b/speechbrain/pretrained/training.py
@@ -5,13 +5,17 @@
 * Artem Ploujnikov 2021
 """
 import os
+import logging
 import shutil
 
+logger = logging.getLogger(__name__)
+
 
 def save_for_pretrained(
     hparams,
     min_key=None,
     max_key=None,
+    ckpt_predicate=None,
     pretrainer_key="pretrainer",
     checkpointer_key="checkpointer",
 ):
@@ -24,12 +28,14 @@ def save_for_pretrained(
     ---------
     hparams: dict
         the hyperparameter file
-    max_key : str
+    max_key: str
         Key to use for finding best checkpoint (higher is better).
         By default, passed to ``self.checkpointer.recover_if_possible()``.
-    min_key : str
+    min_key: str
         Key to use for finding best checkpoint (lower is better).
-        By def  ault, passed to ``self.checkpointer.recover_if_possible()``.
+        By default, passed to ``self.checkpointer.recover_if_possible()``.
+    ckpt_predicate: callable
+        a filter predicate to locate checkpoints
     checkpointer_key: str
         the key under which the checkpointer is stored
     pretrained_key: str
@@ -42,21 +48,42 @@ def save_for_pretrained(
         )
     pretrainer = hparams[pretrainer_key]
     checkpointer = hparams[checkpointer_key]
-    checkpoint = checkpointer.find_checkpoint(min_key=min_key, max_key=max_key)
+    checkpoint = checkpointer.find_checkpoint(
+        min_key=min_key, max_key=max_key, ckpt_predicate=ckpt_predicate
+    )
+    if checkpoint:
+        logger.info(
+            "Saving checkpoint '%s' a pretrained model", checkpoint.path
+        )
+        pretrainer_keys = set(pretrainer.loadables.keys())
+        checkpointer_keys = set(checkpoint.paramfiles.keys())
+        keys_to_save = pretrainer_keys & checkpointer_keys
+        for key in keys_to_save:
+            source_path = checkpoint.paramfiles[key]
+            if not os.path.exists(source_path):
+                raise ValueError(
+                    f"File {source_path} does not exist in the checkpoint"
+                )
+            target_path = pretrainer.paths[key]
+            dirname = os.path.dirname(target_path)
+            if not os.path.exists(dirname):
+                os.makedirs(dirname)
+            if os.path.exists(target_path):
+                os.remove(target_path)
+            shutil.copyfile(source_path, target_path)
+        saved = True
+    else:
+        logger.info(
+            "Unable to find a matching checkpoint for min_key = %s, max_key = %s",
+            min_key,
+            max_key,
+        )
+        checkpoints = checkpointer.list_checkpoints()
+        checkpoints_str = "\n".join(
+            f"{checkpoint.path}: {checkpoint.meta}"
+            for checkpoint in checkpoints
+        )
+        logger.info("Available checkpoints: %s", checkpoints_str)
+        saved = False
 
-    pretrainer_keys = set(pretrainer.loadables.keys())
-    checkpointer_keys = set(checkpoint.paramfiles.keys())
-    keys_to_save = pretrainer_keys & checkpointer_keys
-    for key in keys_to_save:
-        source_path = checkpoint.paramfiles[key]
-        if not os.path.exists(source_path):
-            raise ValueError(
-                f"File {source_path} does not exist in the checkpoint"
-            )
-        target_path = pretrainer.paths[key]
-        dirname = os.path.dirname(target_path)
-        if not os.path.exists(dirname):
-            os.makedirs(dirname)
-        if os.path.exists(target_path):
-            os.remove(target_path)
-        shutil.copyfile(source_path, target_path)
+    return saved
