Repository: rwth-i6/returnn

Issue:
Error when setting attention weights layer as output layer
When using the following network dict: https://gist.github.com/robin-p-schmitt/a63bbfd3870935b78c86328a38fae783, I am getting the following error: 

```
  File "/u/schmitt/experiments/transducer/recipe/i6_experiments/users/schmitt/experiments/swb/transducer/tools/dump_attention_weights.py", line 109, in init
    line: rnn.engine.init_network_from_config(net_dict_post_proc=net_dict_add_losses)
    locals:
      rnn = <global> <module 'returnn.__main__' from '/u/schmitt/src/returnn/returnn/__main__.py'>
      rnn.engine = <global> <returnn.tf.engine.Engine object at 0x7ff5f5a47a00>
      rnn.engine.init_network_from_config = <global> <bound method Engine.init_network_from_config of <returnn.tf.engine.Engine object at 0x7ff5f5a47a00>>
      net_dict_post_proc = <not found>
      net_dict_add_losses = <global> <function net_dict_add_losses at 0x7ff5f9e28c10>
...
  File "/u/schmitt/src/returnn/returnn/tf/layers/base.py", line 1342, in LayerBase.get_losses_initialized
    line: return self.__class__.get_losses(reduce_func=reduce_func, layer=self, **self.kwargs)
    locals:
      self = <local> <RecLayer 'label_model' out_type=Data{[T|'label_ground_truth_masked0:masked:time'[B],B], dtype='int32', sparse_dim=Dim{F'label_ground_truth_masked:set-sparse-dim'(1031)}}>
      self.__class__ = <local> <class 'returnn.tf.layers.rec.RecLayer'>
      self.__class__.get_losses = <local> <bound method RecLayer.get_losses of <class 'returnn.tf.layers.rec.RecLayer'>>
      reduce_func = <local> None
      layer = <not found>
      self.kwargs = <local> {'back_prop': True, 'include_eos': True, 'is_output_layer': True, 'name_scope': 'output/rec', '_network': <TFNetwork '' train=<tf.Tensor 'globals/train_flag:0' shape=() dtype=bool>>, '_name': 'label_model', 'n_out': <class 'returnn.util.basic.NotSpecified'>, 'sources': [<SourceLayer 'data:label_g..., len = 13
  File "/u/schmitt/src/returnn/returnn/tf/layers/rec.py", line 725, in RecLayer.get_losses
    line: for loss in template_layer.layer_class_type.get_losses(reduce_func=reduce_func, **template_layer.kwargs):
    locals:
      loss = <local> None
      template_layer = <local> <_TemplateLayer(LengthLayer)(:template:length) label_model/':dyn-tag-accum:1:att_weights' out_type=Data{[B], dtype='int32', ctx=loop('label_ground_truth_masked0:masked:time'[B])} (construction stack None)>
      template_layer.layer_class_type = <local> <class 'returnn.tf.layers.basic.LengthLayer'>
      template_layer.layer_class_type.get_losses = <local> <bound method LayerBase.get_losses of <class 'returnn.tf.layers.basic.LengthLayer'>>
      reduce_func = <local> None
      template_layer.kwargs = <local> {'axis': Dim{'att_t'[B]{ctx=loop('label_ground_truth_masked0:masked:time'[B])}}, '_network': <TFNetwork '/label_model(rec-subnet)' parent_layer=<RecLayer 'label_model' out_type=Data{[T|'label_ground_truth_masked0:masked:time'[B],B], dtype='int32', sparse_dim=Dim{F'label_ground_truth_masked:set-sp..., len = 6
TypeError: get_losses() missing 1 required positional argument: 'network'
```

The full error log can also be seen in the Gist above.

The error seems to be caused by setting `is_output_layer=True` in the attention weights layer in combination with having `prev:att` as input to the `lm` layer. The relevant layers therefore are:

```
"lm": {
                "class": "rec",
                "from": ["input_embed", "prev:att"],
                "n_out": 1024,
                "name_scope": "lm/rec",
                "unit": "nativelstm2",
            },
"att_query": {
                "activation": None,
                "class": "linear",
                "from": "lm",
                "is_output_layer": False,
                "n_out": 1024,
                "with_bias": False,
            },
"att_energy_in": {
                "class": "combine",
                "from": ["att_ctx", "att_query"],
                "kind": "add",
                "n_out": 1024,
            },
"energy_tanh": {
                "activation": "tanh",
                "class": "activation",
                "from": ["att_energy_in"],
            },
"att_energy0": {
                "activation": None,
                "class": "linear",
                "from": ["energy_tanh"],
                "n_out": 1,
                "name_scope": "energy",
                "with_bias": False,
            },
"att_energy": {
                "class": "reinterpret_data",
                "from": "att_energy0",
                "is_output_layer": False,
                "set_dim_tags": {
                    "f": Dim(
                        kind=Dim.Types.Spatial, description="att_heads", dimension=1
                    )
                },
            },
"att_weights0": {
                "axis": "stag:att_t",
                "class": "softmax_over_spatial",
                "energy_factor": 0.03125,
                "from": "att_energy",
            },
"att_weights": {
                "class": "dropout",
                "dropout": 0.0,
                "dropout_noise_shape": {"*": None},
                "from": "att_weights0",
                "is_output_layer": False,
            },
"att0": {
                "add_var2_if_empty": False,
                "class": "dot",
                "from": ["att_val_split", "att_weights"],
                "reduce": "stag:att_t",
                "var1": "f",
                "var2": None,
            },
"att": {"axes": "except_time", "class": "merge_dims", "from": "att0"},
```

In my script, which I use to initialize my network, I call `rnn.engine.init_network_from_config(net_dict_post_proc=net_dict_add_losses)`. The function `net_dict_add_losses`, despite its name, only sets `net_dict["label_model"]["unit"]["att_weights"]["is_output_layer"] = True`. When looking at the log, one can see that the network construction is working before setting this and only fails in the second pass. Therefore it must have something to do with setting `att_weights` as an output layer.

I was not yet able to reproduce the error with a smaller network but maybe the error can still be found.


Discussion:

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
@@ -2324,7 +2324,7 @@ def get_out_data_from_opts(cls, name, algorithm=None, out_dim=None, **kwargs):
     """
     from tensorflow.python.ops import stateful_random_ops
     algorithm = cls.select_algorithm(algorithm)
-    algo_type = tf.random.Algorithm(algorithm)
+    algo_type = tf.random.Algorithm(algorithm)  # noqa
     if algorithm in cls._state_size_dim_tags_cache:
       algo_state_dim = cls._state_size_dim_tags_cache[algorithm]
     else:
@@ -3351,7 +3351,7 @@ def _get_merge_axes(cls, axes, keep_order, input_data, name):
     :rtype: list[int]
     """
     if keep_order:
-      assert isinstance(axes, (tuple, list, typing.Sequence)), (
+      assert isinstance(axes, (tuple, list, typing.Sequence)) and not isinstance(axes, str), (
         "%s: axes %r must be a list or tuple, to have a well defined order in input %s" % (name, axes, input_data))
       axes_ = []
       for axis in axes:
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
@@ -1683,13 +1683,12 @@ def _add_template_layer(self, layer_name, layer_dict):
     layer_dict = layer_dict.copy()
     layer_class_name = layer_dict.pop("class")
     layer_class = get_layer_class(layer_class_name)
-    layer_dict["_network"] = self.net
-    layer_dict["_name"] = layer_name
+    layer_dict.update(dict(name=layer_name, _name=layer_name, network=self.net, _network=self.net))
     layer_class.transform_config_dict(
       layer_dict, network=self.net, get_layer=lambda _name: self.layer_data_templates[_name])
-    out = layer_class.get_out_data_from_opts(name=layer_name, network=self.net, **layer_dict)
-    out = layer_class.fixup_out_data(output=out, network=self.net, **layer_dict)
-    layer.init(output=out, layer_class=layer_class, **layer_dict)
+    layer_dict["output"] = layer_class.get_out_data_from_opts(**layer_dict)
+    layer_dict["output"] = layer_class.fixup_out_data(**layer_dict)
+    layer.init(layer_class=layer_class, **layer_dict)
     self.layer_data_templates[layer_name] = layer
     return layer
 
diff --git a/returnn/tf/network.py b/returnn/tf/network.py
--- a/returnn/tf/network.py
+++ b/returnn/tf/network.py
@@ -1592,7 +1592,7 @@ def get_fetches_dict(self, config=None, should_train=None, should_eval=None,
     if config is None:
       config = self.get_config()
     if should_train is None:
-      should_train = self.train_flag
+      should_train = self.train_flag is not False
     if should_eval is None:
       should_eval = self.eval_flag
     use_horovod_reduction = False


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -5455,6 +5455,282 @@ def test_reclayer_explicit_rec_ff():
     session.run(network.get_default_output_layer().output.placeholder, feed_dict=make_feed_dict(network.extern_data))
 
 
+def test_reclayer_att_weights_output_layer():
+  # https://github.com/rwth-i6/returnn/issues/1027
+  from returnn.tf.util.data import batch_dim, SpatialDim, FeatureDim
+
+  time_dim = SpatialDim('time')
+  input_dim = FeatureDim('input', 13)
+  align_classes_dim = FeatureDim("align", 5)
+
+  config = Config({"extern_data": {
+    'data': {
+      'dim_tags': (
+        batch_dim,
+        time_dim,
+        input_dim
+      ),
+      'dtype': 'float32',
+      'available_for_inference': True
+    },
+    'alignment': {
+      'dim_tags': (
+        batch_dim,
+        time_dim,
+      ),
+      "sparse_dim": align_classes_dim,
+      'dtype': 'int32',
+      'available_for_inference': True
+    },
+  }})
+
+  net_dict = {
+    "encoder": {"class": "copy", "from": "data"},
+    "existing_alignment": {"class": "copy", "from": "data:alignment"},
+    "is_label": {
+      "class": "compare",
+      "from": "existing_alignment",
+      "kind": "not_equal",
+      "value": align_classes_dim.dimension - 1,
+    },
+    "label_ground_truth_masked": {
+      "class": "reinterpret_data",
+      "enforce_batch_major": True,
+      "from": "label_ground_truth_masked0",
+      "register_as_extern_data": "label_ground_truth",
+      "set_sparse_dim": align_classes_dim - 1,
+    },
+    "label_ground_truth_masked0": {
+      "class": "masked_computation",
+      "from": "existing_alignment",
+      "mask": "is_label",
+      "unit": {"class": "copy", "from": "data"},
+    },
+    "label_model": {
+      "back_prop": True,
+      "class": "rec",
+      "from": "data:label_ground_truth",
+      "include_eos": True,
+      "is_output_layer": True,
+      "name_scope": "output/rec",
+      "unit": {
+        "att": {"axes": ["stag:heads", input_dim], "class": "merge_dims", "from": "att0"},
+        "att0": {
+          "add_var2_if_empty": False,
+          "class": "dot",
+          "from": ["att_val_split", "att_weights"],
+          "reduce": "stag:att_t",
+          "var1": "f",
+          "var2": None,
+        },
+        "att_ctx": {
+          "L2": None,
+          "activation": None,
+          "class": "linear",
+          "dropout": 0.0,
+          "from": "segments",
+          "n_out": 10,
+          "name_scope": "/enc_ctx",
+          "with_bias": True,
+        },
+        "att_energy": {
+          "class": "reinterpret_data",
+          "from": "att_energy0",
+          "is_output_layer": False,
+          "set_dim_tags": {
+            "f": Dim(
+              kind=Dim.Types.Spatial, description="att_heads", dimension=1
+            )
+          },
+        },
+        "att_energy0": {
+          "activation": None,
+          "class": "linear",
+          "from": ["energy_tanh"],
+          "n_out": 1,
+          "name_scope": "energy",
+          "with_bias": False,
+        },
+        "att_energy_in": {
+          "class": "combine",
+          "from": ["att_ctx", "att_query"],
+          "kind": "add",
+          "n_out": 10,
+        },
+        "att_query": {
+          "activation": None,
+          "class": "linear",
+          "from": "lm",
+          "is_output_layer": False,
+          "n_out": 10,
+          "with_bias": False,
+        },
+        "att_val": {"class": "copy", "from": "segments"},
+        "att_val_split": {
+          "class": "reinterpret_data",
+          "from": "att_val_split0",
+          "set_dim_tags": {
+            "dim:1": Dim(
+              kind=Dim.Types.Spatial, description="att_heads", dimension=1
+            )
+          },
+        },
+        "att_val_split0": {
+          "axis": "f",
+          "class": "split_dims",
+          "dims": (1, -1),
+          "from": "att_val",
+        },
+        "att_weights": {
+          "class": "dropout",
+          "dropout": 0.0,
+          "dropout_noise_shape": {"*": None},
+          "from": "att_weights0",
+          "is_output_layer": True,
+        },
+        "att_weights0": {
+          "axis": "stag:att_t",
+          "class": "softmax_over_spatial",
+          "energy_factor": 0.03125,
+          "from": "att_energy",
+        },
+        "energy_tanh": {
+          "activation": "tanh",
+          "class": "activation",
+          "from": ["att_energy_in"],
+        },
+        "input_embed": {
+          "class": "copy",
+          "from": "prev:target_embed",
+        },
+        "lm": {
+          "class": "rec",
+          "from": ["input_embed", "prev:att"],
+          "n_out": 10,
+          "name_scope": "lm/rec",
+          "unit": "nativelstm2",
+        },
+        "output": {
+          "beam_size": 4,
+          "cheating": "exclusive",
+          "class": "choice",
+          "from": "data",
+          "initial_output": 0,
+          "target": "label_ground_truth",
+        },
+        "segment_lens": {
+          "axis": "t",
+          "class": "gather",
+          "from": "base:data:segment_lens_masked",
+          "position": ":i",
+        },
+        "segment_starts": {
+          "axis": "t",
+          "class": "gather",
+          "from": "base:data:segment_starts_masked",
+          "position": ":i",
+        },
+        "segments": {
+          "class": "reinterpret_data",
+          "from": "segments0",
+          "set_dim_tags": {
+            "stag:sliced-time:segments": Dim(
+              kind=Dim.Types.Spatial, description="att_t"
+            )
+          },
+        },
+        "segments0": {
+          "class": "slice_nd",
+          "from": "base:encoder",
+          "size": "segment_lens",
+          "start": "segment_starts",
+        },
+        "target_embed": {
+          "activation": None,
+          "class": "linear",
+          "from": "output",
+          "initial_output": 0,
+          "n_out": 6,
+          "with_bias": False,
+        },
+      },
+    },
+    "output": {
+      "back_prop": False,
+      "class": "rec",
+      "from": "encoder",
+      "include_eos": True,
+      "initial_output": 0,
+      "size_target": None,
+      "target": "alignment",
+      "unit": {
+        "const1": {"class": "constant", "value": 1},
+        "output": {
+          "beam_size": 4,
+          "cheating": None,
+          "class": "choice",
+          "from": "data",
+          "initial_output": 0,
+          "input_type": "log_prob",
+          "length_normalization": False,
+          "target": "alignment",
+        },
+        "output_emit": {
+          "class": "compare",
+          "from": "output",
+          "initial_output": True,
+          "kind": "not_equal",
+          "value": align_classes_dim.dimension - 1,
+        },
+        "segment_lens": {
+          "class": "combine",
+          "from": ["segment_lens0", "const1"],
+          "is_output_layer": True,
+          "kind": "add",
+        },
+        "segment_lens0": {
+          "class": "combine",
+          "from": [":i", "segment_starts"],
+          "kind": "sub",
+        },
+        "segment_starts": {
+          "class": "switch",
+          "condition": "prev:output_emit",
+          "false_from": "prev:segment_starts",
+          "initial_output": 0,
+          "is_output_layer": True,
+          "true_from": ":i",
+        },
+      },
+    },
+    "segment_lens_masked": {
+      "class": "masked_computation",
+      "from": "output/segment_lens",
+      "mask": "is_label",
+      "out_spatial_dim": Dim(kind=Dim.Types.Spatial, description="label-axis"),
+      "register_as_extern_data": "segment_lens_masked",
+      "unit": {"class": "copy", "from": "data"},
+    },
+    "segment_starts_masked": {
+      "class": "masked_computation",
+      "from": "output/segment_starts",
+      "mask": "is_label",
+      "out_spatial_dim": Dim(kind=Dim.Types.Spatial, description="label-axis"),
+      "register_as_extern_data": "segment_starts_masked",
+      "unit": {"class": "copy", "from": "data"},
+    },
+  }
+
+  with make_scope() as session:
+    network = TFNetwork(config=config, train_flag=True)
+    network.construct_from_dict(net_dict)
+    network.initialize_params(session)
+    from test_TFNetworkLayer import make_feed_dict
+    fetches = network.get_fetches_dict()
+    fetches["out"] = network.get_layer("label_model/att_weights").output.placeholder
+    session.run(fetches, feed_dict=make_feed_dict(network.extern_data))
+
+
 def test_convert_lstm_params_save_load():
   """
   Test conversions from different units to different units.