Repository: rwth-i6/returnn

Issue:
Extremely slow network construction (or maybe infinite loop, unclear)
At the same time, memory increases all the time (slowly), so after some minutes, it would crash due to out-of-memory.

This is via a transducer search config generated via returnn-common, [this config](https://gist.github.com/albertz/e16974e8081030934cd27ea3981d7503), but the problem also occurs when I dump it as pure RETURNN net dict, as in [this config](https://gist.github.com/albertz/9dba58ff6119ba63af8d987144d8f4dc/575e7c7ef813227c2622fb513b649cc713ed17d2). Simplified further, [this config](https://gist.github.com/albertz/9dba58ff6119ba63af8d987144d8f4dc).

In the log you see this as some of the last messages:
```
...
layer /lm/lstm/'output': [F'embed'(256),F|F'4*lstm'(4096)] float32
layer /lm/'lstm': [F'embed'(256),F|F'4*lstm'(4096)] float32
layer /lm/'output': [F'embed'(256),F|F'4*lstm'(4096)] float32
layer /'lm': [F'embed'(256),F|F'4*lstm'(4096)] float32
layer /beam_search/'random_4': [F'lstm'(1024),F|F'readout'(1000)] float32
layer /beam_search/'readout_in_lm_weight': [F'lstm'(1024),F|F'readout'(1000)] float32
layer /'readout_in_bias': [F|F'readout'(1000)] float32
layer /beam_search/'random_6': [F'readout//2'(500),F|F'phones'(61)] float32
layer /beam_search/'out_nb_label_logits_weight': [F'readout//2'(500),F|F'phones'(61)] float32
layer /beam_search/'out_nb_label_logits_bias': [F|F'phones'(61)] float32
```

When looking at the stack trace in a debugger, rec template construction is involved.



Discussion:
To get an impression of the (maybe endless, or very long) loop, in `_SubnetworkRecCell._construct_template().RecSubnetGetTemplateLayer.__call__()`, add:
```python
        print(" ", self, "-" * len(ConstructCtx.layers), name)
```
You get some output like:
```
...
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:embed_bias
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:base:lm
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:readout_in_lm_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- model/ne
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ---------------- model/combine_bc
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ---------------- base:base:readout_in_bias
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ---------------- model/combine_bc
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ----------------- model/readout_in_am
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ----------------- model/masked_computation/readout_in_lm
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------ model/masked_computation
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- model/reinterpret_set_sparse_dim
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:embed_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:embed_bias
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:base:lm
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:readout_in_lm_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- model/ne
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- model/reinterpret_set_sparse_dim
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:embed_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:embed_bias
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:base:lm
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:readout_in_lm_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- model/ne
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- model/reinterpret_set_sparse_dim
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:embed_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:embed_bias
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:base:lm
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:readout_in_lm_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- model/ne
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------ model/masked_computation
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------ model/masked_computation
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- model/reinterpret_set_sparse_dim
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:embed_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------------- base:embed_bias
...
```
Repeating a lot...

Maybe there is actually some bug in the config, and some exception occurs, however, due to the template construction logic, it keeps trying.

Maybe the flat net construction (#992) would also solve this.
The masked computation layer is not involved. I simplified it further (updated config is in the main post) and remove that part. Now the debug output:
```
...
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------ base:readout_in_lm_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------ model/out_wb_label_logits/dot
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------ base:out_wb_label_logits_bias
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------ model/out_wb_label_logits/dot
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------- model/reduce_out
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> -------- model/combine_bc
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> --------- model/readout_in_am
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> --------- model/readout_in_lm
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ---------- model/readout_in_lm/output
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ----------- model/readout_in_lm/dot
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------ model/lm
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------ base:readout_in_lm_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------ model/lm
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------ base:readout_in_lm_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------ model/lm
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ------------ base:readout_in_lm_weight
  <_SubnetworkRecCell '/beam_search/loop(rec-subnet)'> ----------- model/readout_in_lm/dot
...
```
I was thinking about a complete redesign of the net construction (#1128). But the main problem here is actually in the rec layer subnet template construction heuristic (for resolving circular dependencies). This is somewhat independent. Separate issue on redesigning the rec layer subnet template construction heuristic: #1129
My assumption is that some (actual correct) exception occurs in some layer (e.g. `model/lm` as you see that often last in the debug output) but it keeps trying again and again.

We should try to detect this somehow. But how?

Resolving patch:
diff --git a/returnn/tf/layers/basic.py b/returnn/tf/layers/basic.py
--- a/returnn/tf/layers/basic.py
+++ b/returnn/tf/layers/basic.py
@@ -157,6 +157,8 @@ def get_concat_sources_data_template(src_layers, out_dim=None, allow_broadcast_a
   dim = None
   common_source = Data.get_common_data(
     [s.output for s in src_layers], ignore_feature_dim=True, allow_broadcast_all_sources=allow_broadcast_all_sources)
+  if not common_source.have_feature_axis():  # e.g. during template construction
+    return common_source
   for layer in src_layers:
     # Note: We do not perform much compatibility checks at this point,
     # as this is for a template only anyway.
diff --git a/returnn/tf/layers/rec.py b/returnn/tf/layers/rec.py
--- a/returnn/tf/layers/rec.py
+++ b/returnn/tf/layers/rec.py
@@ -500,7 +500,7 @@ def get_out_data_from_opts(cls, name, network, sources, unit, axis=None, out_dim
         out.time_dim_axis = out.get_axis_from_description(axis)
       if out.have_time_axis() and axis == out.get_time_dim_tag():
         out = out.copy_as_time_batch_major()
-      else:
+      elif out.have_batch_axis():
         # We expect to be inside another RecLayer, and should do a single step (like RnnCellLayer).
         out = out.copy_as_batch_major()  # The output is then [B,F]
     else:
@@ -1402,14 +1402,23 @@ def add_templated_layer(lself, name, layer_class, **layer_desc):
         layer_desc["name"] = name
         layer_desc["network"] = self.net
         old_layer_kwargs = layer_.kwargs
-        layer_.kwargs = layer_desc.copy()  # set it now already for better debugging
+        # set it now already for better debugging
+        layer_.layer_class_type = layer_class
+        layer_.kwargs = layer_desc.copy()
         if "output" not in layer_.kwargs:
           if old_layer_kwargs and "output" in old_layer_kwargs:
             # First copy old output. Maybe the get_out_data_from_opts raises an exception,
             # and we don't want this to be unset.
             layer_.kwargs["output"] = old_layer_kwargs["output"]
           layer_.kwargs["output"] = layer_class.get_out_data_from_opts(**layer_desc)
-        layer_.kwargs["output"] = layer_class.fixup_out_data(**layer_.kwargs)
+        # noinspection PyBroadException
+        try:
+          layer_.kwargs["output"] = layer_class.fixup_out_data(**layer_.kwargs)
+        except Exception:  # most likely VerifyOutShapeException
+          if lself.got_uninitialized_deps_count > 0:
+            pass  # ignore in this case
+          else:
+            raise
         layer_.kwargs["output"].sanity_check(ignore_placeholder=True)  # placeholder might be overwritten later
         layer_.init(layer_class=layer_class, **layer_.kwargs)
         if layer_.need_last:
@@ -1484,10 +1493,13 @@ def __call__(lself, name, is_prev_time_frame=False):
           # * layer_class.transform_config_dict (via construct_layer)
           # * layer_class.get_out_data_from_opts (via add_templated_layer)
           ConstructCtx.partially_finished.append(layer_)
-        default_get_layer = RecSubnetGetTemplateLayer(parent=lself, parent_name=_name)
+        if lself.allow_uninitialized_template:
+          lself._add_uninitialized_count()
+          return layer_
         default_success = False  # whether construction was successful with default_get_layer
         ConstructCtx.layers.append(layer_)
         try:
+          initial_output_layer = None
           # If not initialized yet and this is still the dummy data template,
           # see if we can figure out a petter template.
           if not layer_.is_initialized and layer_.output.batch_ndim == 0:
@@ -1497,50 +1509,49 @@ def __call__(lself, name, is_prev_time_frame=False):
               if isinstance(initial_output, str) and initial_output.startswith("base:"):
                 initial_output_layer = _parent_get_layer(initial_output[len("base:"):])
                 layer_.output = initial_output_layer.output.copy_template(name="%s_output" % name)
-          # See how far we can get without recursive layer construction.
-          # We only want to get the data template for now.
-          # If that fails in some way,
-          # try another time but only allowing recursive layer construction for the first get_layer call.
-          # E.g. the CombineLayer and some other layers determine the output format via the first source.
-          # Also, first try without allowing to access uninitialized templates,
-          # as they might propagate wrong Data format info (they have a dummy Data format set).
-          # Only as a last resort, allow this.
-          get_layer_candidates = []  # type: typing.List[RecSubnetGetTemplateLayer]
-          # noinspection PyProtectedMember
-          if lself.iterative_testing and name not in self.net._construction_stack.layers:
-            # We can get away with only two variants, because the reconstruction code below
-            # for partially finished layers will make sure that everything is correct.
-            get_layer_candidates = [
-              default_get_layer,
-              RecSubnetGetTemplateLayer(allow_uninitialized_template=True, parent=lself, parent_name=_name),
-            ]
-          for get_layer in get_layer_candidates:
+
+          ConstructCtx.most_recent = list(ConstructCtx.layers)
+          default_get_layer = RecSubnetGetTemplateLayer(parent=lself, parent_name=_name)
+          last_get_layer = default_get_layer
+          # noinspection PyBroadException
+          try:
+            res_layer = self.net.construct_layer(
+              net_dict=self.net_dict, name=name,
+              get_layer=default_get_layer, add_layer=default_get_layer.add_templated_layer)
+            default_success = True
+          except Exception:
+            ConstructCtx.collect_exception(layer_name=name)
+            if layer_.is_initialized:
+              # Return anyway. This will be resolved later.
+              # Also, this is important to not keep trying again and again potentially for a very long time.
+              # https://github.com/rwth-i6/returnn/issues/1127
+              pass
+            elif lself.iterative_testing or initial_output_layer:
+              pass  # cover this below
+            else:
+              raise
+
+          if not layer_.is_initialized and lself.iterative_testing:
+            shallow_get_layer = RecSubnetGetTemplateLayer(
+              allow_uninitialized_template=True, parent=lself, parent_name=_name)
+            last_get_layer = shallow_get_layer
             # noinspection PyBroadException
             try:
-              res_layer = self.net.construct_layer(
+              self.net.construct_layer(
                 net_dict=self.net_dict, name=name,
-                get_layer=get_layer, add_layer=get_layer.add_templated_layer)
-              if get_layer is default_get_layer:
-                default_success = True
-              break  # we did it, so get out of the loop
+                get_layer=shallow_get_layer, add_layer=shallow_get_layer.add_templated_layer)
             except Exception:
               ConstructCtx.collect_exception(layer_name=name)
-          if not default_success:
-            # Now, do again, but with full recursive layer construction, to determine the dependencies.
-            ConstructCtx.most_recent = list(ConstructCtx.layers)
-            try:
-              default_get_layer.reset()
-              res_layer = self.net.construct_layer(
-                net_dict=self.net_dict, name=name,
-                get_layer=default_get_layer, add_layer=default_get_layer.add_templated_layer)
-              default_success = True
-            except NetworkConstructionDependencyLoopException:
-              if layer_.is_initialized and lself.iterative_testing and not lself.reconstruct:
-                pass  # Return anyway. This will be resolved later.
-              else:
+              if not initial_output_layer:
                 raise
-            except Exception:
-              raise
+
+          if initial_output_layer and layer_ in ConstructCtx.partially_finished:
+            # We can trust the initial output, and even ignore whatever else we got from the construct_layer above.
+            # The construct_layer is just to setup the deps and kwargs.
+            assert layer_.kwargs
+            layer_.kwargs["output"] = initial_output_layer.output.copy_template(name="%s_output" % name)
+            last_get_layer.add_templated_layer(
+              layer_class=layer_.layer_class_type, **layer_.kwargs)
 
           if res_layer and res_layer is not layer_:
             assert isinstance(res_layer, _TemplateLayer), "unexpected %s, via %s" % (res_layer, layer_)
@@ -1553,10 +1564,11 @@ def __call__(lself, name, is_prev_time_frame=False):
           if exc.network is self.net and exc.layer_name == name:
             assert not layer_.is_initialized
             # Clean up any references.
-            del self.layer_data_templates[name]
+            layer_ = self.layer_data_templates.pop(name)
             ConstructCtx.partially_finished.remove(layer_)
-            if len(ConstructCtx.layers) >= 2:
-              ConstructCtx.layers[-2].remove_dependency(layer_, is_prev_time_frame=is_prev_time_frame)
+            for other in list(layer_.used_by):
+              if isinstance(other, _TemplateLayer):
+                other.remove_dependency(layer_, is_prev_time_frame=is_prev_time_frame)
           raise
 
         finally:
@@ -3886,6 +3898,7 @@ def __init__(self, network, name, construct_stack=None, cell=None):
     self.dependencies = []  # type: typing.List[LayerBase]
     self.cur_frame_dependencies = []  # type: typing.List[LayerBase]
     self.prev_frame_dependencies = []  # type: typing.List[_TemplateLayer]
+    self.used_by = []  # type: typing.List[LayerBase]  # reverse deps
     self.construct_stack = construct_stack
     self._template_base = None  # type: typing.Optional[_TemplateLayer]
     self._cell = cell
@@ -4071,6 +4084,8 @@ def add_dependency(self, layer, is_prev_time_frame):
     """
     if layer not in self.dependencies:
       self.dependencies.append(layer)
+      if isinstance(layer, _TemplateLayer):
+        layer.used_by.append(self)
     if is_prev_time_frame:
       assert isinstance(layer, _TemplateLayer)
       if layer not in self.prev_frame_dependencies:
@@ -4085,6 +4100,8 @@ def remove_dependency(self, layer, is_prev_time_frame):
     :param bool is_prev_time_frame:
     """
     self.dependencies.remove(layer)
+    if isinstance(layer, _TemplateLayer):
+      layer.used_by.remove(self)
     if is_prev_time_frame:
       assert isinstance(layer, _TemplateLayer)
       self.prev_frame_dependencies.remove(layer)
@@ -4527,15 +4544,16 @@ def get_out_data_from_opts(cls, n_out, name, sources=(), **kwargs):
     :rtype: Data
     """
     sources_data = Data.get_common_data([src.output for src in sources if src], ignore_feature_dim=True)
-    feat = FeatureDim("%s:rnn_cell_feat" % name, dimension=n_out, auto_generated=True)
+    batch_dim = sources_data.get_batch_dim_tag() if sources_data and sources_data.have_batch_axis() else None
+    time_dim_axis, batch_dim_axis = None, None
+    dim_tags = []
     if sources_data and sources_data.have_time_axis():
-      dim_tags = (sources_data.get_time_dim_tag(), sources_data.get_batch_dim_tag(), feat)
-      batch_dim_axis = 1
-      time_dim_axis = 0
-    else:
-      dim_tags = (sources_data.get_batch_dim_tag(), feat)
-      batch_dim_axis = 0
-      time_dim_axis = None
+      time_dim_axis = len(dim_tags)
+      dim_tags.append(sources_data.get_time_dim_tag())
+    if batch_dim:
+      batch_dim_axis = len(dim_tags)
+      dim_tags.append(batch_dim)
+    dim_tags.append(FeatureDim("%s:rnn_cell_feat" % name, dimension=n_out, auto_generated=True))
     return Data(
       name="%s_output" % name,
       dim_tags=dim_tags, dim=n_out,
@@ -8147,7 +8165,7 @@ def get_out_data_from_opts(cls, network, masked_from=None, in_spatial_dim=None,
       output = output.copy_template_replace_dim_tag(axis=axis, new_dim_tag=out_spatial_dim)
     inside_rec_time_dim = network.get_inside_rec_time_dim(inside_loop=True)
     over_rec_time_dim = network.get_inside_rec_time_dim(inside_loop=False)
-    if network.is_inside_rec_layer():
+    if network.is_inside_rec_layer() and output.have_batch_axis():
       output = output.copy_as_batch_major()
     elif not masked_from and in_spatial_dim == over_rec_time_dim and over_rec_time_dim and not inside_rec_time_dim:
       # We will automatically unmask again (https://github.com/rwth-i6/returnn/pull/976).


Test patch: