returnn 1.20240307.132239__tar.gz → 1.20240309.14452__tar.gz

{returnn-1.20240307.132239/returnn.egg-info → returnn-1.20240309.14452}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240307.132239
+Version: 1.20240309.14452
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20240309.14452/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240309.014452'
+long_version = '1.20240309.014452+git.e7222af'

{returnn-1.20240307.132239 → returnn-1.20240309.14452}/returnn/datasets/util/vocabulary.py

@@ -216,6 +216,8 @@ class Vocabulary(object):
         if isinstance(label, str):
             return self.label_to_id(label, default=default)
         if isinstance(label, int):
+            if label == -1 and allow_none:
+                return None
             if self.is_id_valid(label):
                 return label
             if default is KeyError:

{returnn-1.20240307.132239 → returnn-1.20240309.14452}/returnn/engine/base.py

@@ -130,9 +130,9 @@ class EngineBase:
 
         load_model_epoch_filename = util.get_checkpoint_filepattern(config.value("load", ""))
         if load_model_epoch_filename:
-            assert os.path.exists(
-                load_model_epoch_filename + util.get_model_filename_postfix()
-            ), "load option %r, file %r does not exist" % (
+            assert os.path.exists(load_model_epoch_filename + util.get_model_filename_postfix()) or os.path.exists(
+                load_model_epoch_filename
+            ), "config option load=%r, file %r does not exist" % (
                 config.value("load", ""),
                 load_model_epoch_filename + util.get_model_filename_postfix(),
             )

{returnn-1.20240307.132239 → returnn-1.20240309.14452}/returnn/torch/engine.py

@@ -666,11 +666,13 @@ class Engine(EngineBase):
         checkpoint_state = None
         if model_epoch_filename:
             filename = model_epoch_filename + util.get_model_filename_postfix()
+            if not os.path.exists(filename) and os.path.exists(model_epoch_filename):
+                filename = model_epoch_filename
             print("Load model %s" % (filename,), file=log.v4)
             checkpoint_state = torch.load(filename, map_location=self._device)
             if epoch is None:
-                epoch = checkpoint_state["epoch"]
-            step = checkpoint_state["step"]
+                epoch = checkpoint_state.get("epoch", self._start_epoch or 1)
+            step = checkpoint_state.get("step", 1)
             print(f"  epoch {epoch}, global train step {step}", file=log.v4)
             # The checkpoint was saved when the step was already increased (but not the epoch yet).
             # Restore the last step.
@@ -690,9 +692,10 @@ class Engine(EngineBase):
         missing_keys = set()
         unexpected_keys = set()
         if checkpoint_state is not None:
-            loaded_state_keys.update(checkpoint_state["model"].keys())
+            model_state = checkpoint_state.get("model", checkpoint_state)
+            loaded_state_keys.update(model_state.keys())
             missing_keys_main_ckpt, unexpected_keys_main_ckpt = self._pt_model.load_state_dict(
-                checkpoint_state["model"], strict=False
+                model_state, strict=False
             )
             missing_keys.update(missing_keys_main_ckpt)
             unexpected_keys.update(unexpected_keys_main_ckpt)
@@ -702,6 +705,10 @@ class Engine(EngineBase):
                     + ", ".join(map(repr, sorted(unexpected_keys_main_ckpt))),
                     file=log.v4,
                 )
+            del model_state
+        # https://github.com/rwth-i6/returnn/issues/1345
+        del checkpoint_state
+        gc.collect()
 
         preload_from_files = self.config.typed_value("preload_from_files", {})
         if preload_from_files:
@@ -788,9 +795,6 @@ class Engine(EngineBase):
                 )
             )
 
-        # https://github.com/rwth-i6/returnn/issues/1345
-        del checkpoint_state
-        gc.collect()
         self._pt_model.to(self._device)
 
         if model_epoch_filename and is_training:

{returnn-1.20240307.132239 → returnn-1.20240309.14452/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240307.132239
+Version: 1.20240309.14452
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20240307.132239/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240307.132239'
-long_version = '1.20240307.132239+git.c1ac801'