add GetBestPtCheckpointJob (#432)

returnn/training.py

@@ -3,6 +3,7 @@
     "Checkpoint",
     "GetBestEpochJob",
     "GetBestTFCheckpointJob",
+    "GetBestPtCheckpointJob",
     "PtCheckpoint",
     "ReturnnModel",
     "ReturnnTrainingFromFileJob",
@@ -797,6 +798,40 @@ def run(self):
         )
 
 
+class GetBestPtCheckpointJob(GetBestEpochJob):
+    """
+    Analog to GetBestTFCheckpointJob, just for torch checkpoints.
+    """
+
+    def __init__(self, model_dir: tk.Path, learning_rates: tk.Path, key: str, index: int = 0):
+        """
+
+        :param Path model_dir: model_dir output from a ReturnnTrainingJob
+        :param Path learning_rates: learning_rates output from a ReturnnTrainingJob
+        :param str key: a key from the learning rate file that is used to sort the models
+            e.g. "dev_score_output/output_prob"
+        :param int index: index of the sorted list to access, 0 for the lowest, -1 for the highest score
+        """
+        super().__init__(model_dir, learning_rates, key, index)
+        self.out_checkpoint = PtCheckpoint(self.output_path("checkpoint.pt"))
+
+    def run(self):
+        super().run()
+
+        try:
+            os.link(
+                os.path.join(self.model_dir.get_path(), "epoch.%.3d.pt" % self.out_epoch.get()),
+                self.out_checkpoint.path,
+            )
+        except OSError:
+            # the hardlink will fail when there was an imported job on a different filesystem,
+            # thus do a copy instead then
+            shutil.copy(
+                os.path.join(self.model_dir.get_path(), "epoch.%.3d.pt" % self.out_epoch.get()),
+                self.out_checkpoint.path,
+            )
+
+
 class AverageTFCheckpointsJob(Job):
     """
     Compute the average of multiple specified Tensorflow checkpoints using the tf_avg_checkpoints script from Returnn