returnn 1.20231214.165403__tar.gz → 1.20231217.143104__tar.gz

{returnn-1.20231214.165403/returnn.egg-info → returnn-1.20231217.143104}/PKG-INFO

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

returnn-1.20231217.143104/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20231217.143104'
+long_version = '1.20231217.143104+git.7a71805'

{returnn-1.20231214.165403 → returnn-1.20231217.143104}/returnn/torch/engine.py

@@ -116,6 +116,7 @@ class Engine(EngineBase):
             torch.cuda.set_device(self._device)
 
         self._log_memory_usage = config.bool("torch_log_memory_usage", False)
+        self._reset_dev_memory_caches = config.bool("reset_dev_memory_caches", False)
 
         amp_options = self.config.opt_typed_value("torch_amp")
         grad_scaler_opts = self.config.typed_value("grad_scaler", NotSpecified)
@@ -248,6 +249,12 @@ class Engine(EngineBase):
             }
         )
 
+    def _maybe_reset_dev_memory_caches(self):
+        if not self._reset_dev_memory_caches:
+            return
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def _reset_dev_memory_stats(self):
         dev = torch.device(self._device)
         if dev.type == "cuda":
@@ -291,6 +298,7 @@ class Engine(EngineBase):
         elapsed_computation_time = 0
 
         self._pt_model.train()
+        self._maybe_reset_dev_memory_caches()
         self._reset_dev_memory_stats()
 
         if self.config.bool("debug_shell_before_train_loop", False):
@@ -413,6 +421,7 @@ class Engine(EngineBase):
         Runs model on all eval datasets and calculates the loss.
         """
         self._pt_model.eval()
+        self._maybe_reset_dev_memory_caches()
         self._reset_dev_memory_stats()
 
         eval_dump_str = []
@@ -807,6 +816,7 @@ class Engine(EngineBase):
         elapsed_computation_time = 0.0
 
         self._pt_model.eval()
+        self._maybe_reset_dev_memory_caches()
         self._reset_dev_memory_stats()
 
         if dataset.supports_seq_order_sorting():

{returnn-1.20231214.165403 → returnn-1.20231217.143104}/returnn/util/basic.py

@@ -645,7 +645,7 @@ def model_epoch_from_filename(filename):
     """
     :param str filename:
     :return: epoch number
-    :rtype: int
+    :rtype: int|None
     """
     # We could check via:
     # tf.contrib.framework.python.framework.checkpoint_utils.load_variable()
@@ -654,16 +654,17 @@ def model_epoch_from_filename(filename):
     # We don't have it in the model, though.
     # For now, just parse it from filename.
     # If TF, and symlink, resolve until no symlink anymore (e.g. if we symlinked the best epoch).
-    while True:
-        tf_meta_fn = "%s.meta" % filename
-        if os.path.exists(tf_meta_fn) and os.path.islink(tf_meta_fn):
-            tf_meta_fn_ = os.readlink(tf_meta_fn)
-            assert tf_meta_fn_.endswith(".meta"), "strange? %s, %s" % (filename, tf_meta_fn)
-            filename = tf_meta_fn_[: -len(".meta")]
-        else:
-            break
+    for potential_ext in [".meta", ".pt"]:
+        if not os.path.exists(filename + potential_ext):
+            continue
+        while os.path.exists(filename + potential_ext) and os.path.islink(filename + potential_ext):
+            fn_with_ext_ = os.readlink(filename + potential_ext)
+            assert fn_with_ext_.endswith(potential_ext), "strange? %s, %s" % (filename, potential_ext)
+            filename = fn_with_ext_[: -len(potential_ext)]
+        break
     m = re.match(".*\\.([0-9]+)", filename)
-    assert m, "no match for %r" % filename
+    if not m:
+        return None
     return int(m.groups()[0])
 
 

{returnn-1.20231214.165403 → returnn-1.20231217.143104}/returnn/util/debug.py

@@ -2,6 +2,8 @@
 Some generic debugging utilities.
 """
 
+from __future__ import annotations
+from typing import Optional, Any, Dict
 import os
 import sys
 import signal
@@ -497,15 +499,19 @@ def init_cuda_not_in_main_proc_check():
     cuda.use.device_number = None
 
 
-def debug_shell(user_ns=None, user_global_ns=None, exit_afterwards=True):
+def debug_shell(
+    user_ns: Optional[Dict[str, Any]] = None,
+    user_global_ns: Optional[Dict[str, Any]] = None,
+    exit_afterwards: bool = True,
+):
     """
     Provides some interactive Python shell.
     Uses IPython if possible.
     Wraps to ``better_exchook.debug_shell``.
 
-    :param dict[str]|None user_ns:
-    :param dict[str]|None user_global_ns:
-    :param bool exit_afterwards: will do sys.exit(1) at the end
+    :param user_ns:
+    :param user_global_ns:
+    :param exit_afterwards: will do sys.exit(1) at the end
     """
     print("Debug shell:")
     from returnn.util.basic import ObjAsDict

{returnn-1.20231214.165403 → returnn-1.20231217.143104/returnn.egg-info}/PKG-INFO

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

returnn-1.20231214.165403/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20231214.165403'
-long_version = '1.20231214.165403+git.cfffc02'