returnn 1.20240103.3709__tar.gz → 1.20240103.123136__tar.gz

{returnn-1.20240103.3709/returnn.egg-info → returnn-1.20240103.123136}/PKG-INFO

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

returnn-1.20240103.123136/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240103.123136'
+long_version = '1.20240103.123136+git.918958c'

{returnn-1.20240103.3709 → returnn-1.20240103.123136}/returnn/torch/engine.py

@@ -230,6 +230,7 @@ class Engine(EngineBase):
 
             epoch_start_global_train_step = self.global_train_step
             epoch_num_tries = 0
+            epoch = self.epoch
             while True:
                 epoch_num_tries += 1
                 if epoch_num_tries > 1:
@@ -237,26 +238,67 @@ class Engine(EngineBase):
                 self.init_train_epoch()
                 try:
                     self.train_epoch()
-                except torch.cuda.OutOfMemoryError as exc:
-                    print("torch.cuda.OutOfMemoryError:", exc, file=log.v1)
-                    # Some heuristics when to retry.
-                    if (
-                        epoch_num_tries <= 3  # heuristic
-                        and self.epoch > start_epoch  # heuristic
-                        and self.global_train_step > epoch_start_global_train_step  # heuristic
-                        and self._save_model_epoch_interval == 1  # cannot load most recent epoch
-                        and not self._torch_distributed_ctx  # need to fix correct sync, all workers need to restart ep
+                except Exception as exc:
+                    if self._handle_train_exception(
+                        exc,
+                        epoch_num_tries=epoch_num_tries,
+                        epoch_start_global_train_step=epoch_start_global_train_step,
+                        start_epoch=start_epoch,
                     ):
-                        print("Retry after OOM", file=log.v3)
                         self._maybe_reset_dev_memory_caches(force=True)
                         self._reset_train_epoch()
-                        assert self.global_train_step == epoch_start_global_train_step
+                        assert self.epoch == epoch and self.global_train_step == epoch_start_global_train_step
                         continue
                     raise
                 break
 
         print(f"Finished training at epoch {self.epoch}, global train step {self.global_train_step}", file=log.v3)
 
+    def _handle_train_exception(
+        self, exc: Exception, *, epoch_num_tries: int, epoch_start_global_train_step: int, start_epoch: int
+    ) -> bool:
+        """
+        :return: True if we should retry, False if reraise
+        """
+        from returnn.util.better_exchook import get_func_from_code_object
+
+        print(f"{type(exc).__name__}: {exc}", file=log.v1)
+
+        # Extend exception message by module call stack.
+        module_names_by_id = {}  # id -> name
+        for name, mod in self._orig_model.named_modules():
+            if id(mod) not in module_names_by_id:
+                module_names_by_id[id(mod)] = name or "(root)"
+        tb = exc.__traceback__
+        exc_ext = ["Module call stack:"]
+        while tb:
+            frame_self = tb.tb_frame.f_locals.get("self")
+            if isinstance(frame_self, (torch.nn.Module, rf.Module)):
+                func = get_func_from_code_object(tb.tb_frame.f_code, frame=tb.tb_frame)
+                if func and func.__name__ and func.__name__.startswith("_"):
+                    tb = tb.tb_next
+                    continue
+                func_name = (func and func.__qualname__) or type(frame_self).__name__
+                exc_ext.append(f"({func_name}) {module_names_by_id.get(id(frame_self), '(unknown)')}")
+            tb = tb.tb_next
+        if len(exc.args) == 1 and isinstance(exc.args[0], str):
+            exc.args = ("\n".join([exc.args[0]] + exc_ext),)
+        else:
+            print("\n".join(exc_ext), file=log.v3)
+
+        if isinstance(exc, torch.cuda.OutOfMemoryError):
+            # Some heuristics when to retry.
+            if (
+                epoch_num_tries <= 3  # heuristic
+                and self.epoch > start_epoch  # heuristic
+                and self.global_train_step > epoch_start_global_train_step  # heuristic
+                and self._save_model_epoch_interval == 1  # cannot load most recent epoch
+                and not self._torch_distributed_ctx  # need to fix correct sync, all workers need to restart ep
+            ):
+                print("Retry after OOM", file=log.v3)
+                return True
+        return False
+
     def init_train_epoch(self):
         """
         init train (sub)epoch. LR etc

{returnn-1.20240103.3709 → returnn-1.20240103.123136/returnn.egg-info}/PKG-INFO

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

{returnn-1.20240103.3709 → returnn-1.20240103.123136}/tests/test_torch_engine.py

@@ -438,6 +438,84 @@ def test_optimizer_convert_aux_param():
         updater_param.load_optimizer(tmp_dir + "/model_buf.opt.pt")
 
 
+class _DemoException(Exception):
+    pass
+
+
+class _TestTorchSubModelRaisingException(torch.nn.Module):
+    def __init__(self, in_features: int, out_features: int):
+        super().__init__()
+        self.lin = torch.nn.Linear(in_features, out_features)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        :param x: [B,T,D]
+        :return: [B,T,D']
+        """
+        x = self.lin(x)
+        if int("1") == 1:
+            raise _DemoException("uh")
+        return x
+
+
+def test_torch_engine_train_exception():
+    class _Model(torch.nn.Module):
+        def __init__(self, **_kwargs):
+            super(_Model, self).__init__()
+            self.sub = _TestTorchSubModelRaisingException(9, 2)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            """
+            :param x: [B,T,D]
+            :return: [B,T,D']
+            """
+            x = self.sub(x)
+            return torch.nn.functional.log_softmax(x, dim=-1)
+
+        @classmethod
+        def train_step(cls, *, model: _Model, extern_data: TensorDict, **_kwargs):
+            """train step"""
+            data: Tensor = extern_data["data"]
+            logits = model(data.raw_tensor)
+            logits_packed = torch.nn.utils.rnn.pack_padded_sequence(
+                logits, data.dims[1].dyn_size_ext.raw_tensor, batch_first=True, enforce_sorted=False
+            )
+            targets = extern_data["classes"]
+            targets_packed = torch.nn.utils.rnn.pack_padded_sequence(
+                targets.raw_tensor, data.dims[1].dyn_size_ext.raw_tensor, batch_first=True, enforce_sorted=False
+            )
+            loss = torch.nn.CrossEntropyLoss(reduction="none")(logits_packed.data, targets_packed.data.long())
+            rf.get_run_ctx().mark_as_loss(name="ce", loss=loss)
+            frame_error = torch.argmax(logits_packed.data, dim=-1).not_equal(targets_packed.data)
+            rf.get_run_ctx().mark_as_loss(name="fer", loss=frame_error, as_error=True)
+
+    config = Config(
+        dict(
+            task="train",
+            device="cpu",
+            extern_data={"data": {"dim": 9}, "classes": {"dim": 2, "sparse": True}},
+            get_model=_Model,
+            train_step=_Model.train_step,
+            batch_size=500,
+            optimizer={"class": "adam"},
+        )
+    )
+    dataset = init_dataset({"class": "Task12AXDataset", "num_seqs": 100, "name": "train"})
+    dataset.init_seq_order(epoch=1)
+
+    with global_config_ctx(config):
+        engine = Engine(config=config)
+        engine.init_train_from_config(train_data=dataset)
+        try:
+            engine.train()
+        except _DemoException as exc:
+            print("got demo exception:", exc)
+            exc_lines = str(exc).splitlines()
+            assert "Module call stack:" in exc_lines and "(_TestTorchSubModelRaisingException.forward) sub" in exc_lines
+        else:
+            raise Exception("did not get expected exception")
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

returnn-1.20240103.3709/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240103.003709'
-long_version = '1.20240103.003709+git.e791ac6'