returnn 1.20240307.22713__tar.gz → 1.20240309.14452__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240307.22713
+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.22713 → 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.22713 → 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.22713 → 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.
@@ -686,25 +688,28 @@ class Engine(EngineBase):
 
         self._create_model(epoch=epoch, step=step)
 
+        loaded_state_keys = set()
+        missing_keys = set()
+        unexpected_keys = set()
         if checkpoint_state is not None:
-            missing_keys, unexpected_keys = self._pt_model.load_state_dict(checkpoint_state["model"], strict=False)
-            if missing_keys:
-                raise Exception(
-                    "\n".join(
-                        [
-                            f"While loading model {filename}:",
-                            "Unexpected key(s) in state_dict: " + ", ".join(map(repr, unexpected_keys)),
-                            "Missing key(s) in state_dict: " + ", ".join(map(repr, missing_keys)),
-                            "Any missing key is an error.",
-                        ]
-                    )
-                )
-            if unexpected_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(
+                model_state, strict=False
+            )
+            missing_keys.update(missing_keys_main_ckpt)
+            unexpected_keys.update(unexpected_keys_main_ckpt)
+            if unexpected_keys_main_ckpt:
                 print(
                     f"Note: While loading {filename}, unexpected key(s) in state_dict: "
-                    + ", ".join(map(repr, unexpected_keys)),
+                    + ", ".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:
             # see `preload_from_files` in tf engine and `returnn.tf.network.CustomCheckpointLoader`
@@ -748,16 +753,48 @@ class Engine(EngineBase):
                         preload_model_state.pop(key)
                 for new_name, name_in_checkpoint in opts.get("var_name_mapping", {}).items():
                     preload_model_state[new_name] = preload_model_state.pop(name_in_checkpoint)
-                missing_keys, _ = self._pt_model.load_state_dict(preload_model_state, strict=False)
-                if missing_keys and not opts.get("ignore_missing", False):
+                missing_keys_preload, unexpected_keys_preload = self._pt_model.load_state_dict(
+                    preload_model_state, strict=False
+                )
+                loaded_state_keys.update(preload_model_state.keys())
+                missing_keys.difference_update(preload_model_state.keys())
+                del preload_model_state
+                gc.collect()
+
+                if opts.get("prefix", ""):
                     prefix_keys = [key for key in self._pt_model.state_dict() if key.startswith(opts.get("prefix", ""))]
-                    missing_prefix_keys = set(prefix_keys).intersection(set(missing_keys))
-                    assert not missing_prefix_keys, f"Missing keys and ignore_missing=False: {missing_prefix_keys}"
-                print(f"Missing keys: {missing_keys}", file=log.v4)
+                else:
+                    prefix_keys = self._pt_model.state_dict().keys()
+                missing_keys_preload = (
+                    set(prefix_keys).intersection(set(missing_keys_preload)).difference(loaded_state_keys)
+                )
+                unexpected_keys_preload = (
+                    set(prefix_keys).intersection(set(unexpected_keys_preload)).difference(loaded_state_keys)
+                )
+                if missing_keys_preload and not opts.get("ignore_missing", False):
+                    missing_keys.update(missing_keys_preload)
+                if missing_keys_preload:
+                    print(f"Missing keys: {missing_keys_preload}", file=log.v4)
+                if unexpected_keys_preload:
+                    print(
+                        f"Note: While loading preload_from_files {opts['filename']}, unexpected key(s) in state_dict: "
+                        + ", ".join(map(repr, sorted(unexpected_keys_preload))),
+                        file=log.v4,
+                    )
+                    unexpected_keys.update(unexpected_keys_preload)
+
+        if missing_keys:
+            raise Exception(
+                "\n".join(
+                    [
+                        f"While loading model {filename}, preload_from_files {bool(preload_from_files)}:",
+                        "Unexpected key(s) in state_dict: " + ", ".join(map(repr, sorted(unexpected_keys))),
+                        "Missing key(s) in state_dict: " + ", ".join(map(repr, sorted(missing_keys))),
+                        "Any missing key is an error.",
+                    ]
+                )
+            )
 
-        # 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.22713 → returnn-1.20240309.14452/returnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20240307.22713
+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.22713/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240307.022713'
-long_version = '1.20240307.022713+git.28576ec'