Repository: rwth-i6/returnn

Issue:
Constraints on layers that are moved out of the RecLayer are lost
Because layers that are moved out of the `RecLayer` are moved either to the `input_layers_net` or the `output_layers_net`, they are not checked for possible constraints in [`get_constraints_value()`](https://github.com/rwth-i6/returnn/blob/master/TFNetworkRecLayer.py#L426). Therefore, the constraints are lost.

I've created a PR with a possible fix as well as test cases for all three situations (no constraints, constraints on moved out input/output layer).


Discussion:

Resolving patch:
diff --git a/TFNetworkRecLayer.py b/TFNetworkRecLayer.py
--- a/TFNetworkRecLayer.py
+++ b/TFNetworkRecLayer.py
@@ -427,7 +427,12 @@ def get_constraints_value(self):
     v = super(RecLayer, self).get_constraints_value()
     from TFUtil import optional_add
     if isinstance(self.cell, _SubnetworkRecCell):
-      for layer in self.cell.net.layers.values():
+      layers = list(self.cell.net.layers.values())
+      if self.cell.input_layers_net:
+        layers += list(self.cell.input_layers_net.layers.values())
+      if self.cell.output_layers_net:
+        layers += list(self.cell.output_layers_net.layers.values())
+      for layer in layers:
         v = optional_add(v, layer.get_constraints_value())
     return v
 


Test patch:
diff --git a/tests/test_TFNetworkRecLayer.py b/tests/test_TFNetworkRecLayer.py
--- a/tests/test_TFNetworkRecLayer.py
+++ b/tests/test_TFNetworkRecLayer.py
@@ -9,7 +9,7 @@
 import sys
 import os
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)) + "/..")
-from nose.tools import assert_equal, assert_is_instance
+from nose.tools import assert_equal, assert_not_equal, assert_is_instance
 from numpy.testing.utils import assert_almost_equal, assert_allclose
 import unittest
 import numpy.testing
@@ -1430,6 +1430,113 @@ def train(net):
   train(train_not_optim_net)
 
 
+def test_rec_layer_move_out_of_loop_keep_constraints():
+  from TFNetworkRecLayer import _SubnetworkRecCell
+  from TFUtil import get_global_train_flag_placeholder
+  n_src_dim = 5
+  n_tgt_dim = 7
+  beam_size = 12
+  config = Config()
+  config.update({"debug_print_layer_output_template": True})
+
+  def get_net_dict(l2_target_embed=0.0, l2_readout_in=0.0):
+    return {
+      "source_embed": {"class": "linear", "activation": None, "with_bias": False, "n_out": 6},
+
+      "lstm0_fw": {"class": "rec", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
+                   "initial_state": "var", "n_out": 10, "direction": 1, "from": ["source_embed"]},
+      "lstm0_bw": {"class": "rec", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
+                   "initial_state": "var", "n_out": 10, "direction": -1, "from": ["source_embed"]},
+
+      "lstm1_fw": {"class": "rec", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
+                   "initial_state": "var", "n_out": 10, "direction": 1, "from": ["lstm0_fw", "lstm0_bw"]},
+      "lstm1_bw": {"class": "rec", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
+                   "initial_state": "var", "n_out": 10, "direction": -1, "from": ["lstm0_fw", "lstm0_bw"]},
+
+      "encoder": {"class": "copy", "from": ["lstm1_fw", "lstm1_bw"]},
+      "enc_ctx": {"class": "linear", "activation": None, "with_bias": True, "from": ["encoder"], "n_out": 10},
+      "fertility": {"class": "linear", "activation": "sigmoid", "with_bias": False, "from": ["encoder"], "n_out": 1},
+
+      "output": {"class": "rec", "from": [], "unit": {
+        'output': {'class': 'choice', 'target': 'classes', 'beam_size': beam_size, 'from': ["output_prob"],
+                   "initial_output": 0},
+        "end": {"class": "compare", "from": ["output"], "value": 0},
+        'target_embed': {'class': 'linear', 'activation': None, "with_bias": False, 'from': ['output'], "n_out": 6,
+                         "initial_output": "apply(0)", "L2": l2_target_embed},
+        "weight_feedback": {"class": "linear", "activation": None, "with_bias": False, "from": ["prev:accum_att_weights"],
+                            "n_out": 10},
+        "prev_s_state": {"class": "get_last_hidden_state", "from": ["prev:s"], "n_out": 20},
+        "prev_s_transformed": {"class": "linear", "activation": None, "with_bias": False, "from": ["prev_s_state"],
+                               "n_out": 10},
+        "energy_in": {"class": "combine", "kind": "add",
+                      "from": ["base:enc_ctx", "weight_feedback", "prev_s_transformed"], "n_out": 10},
+        "energy_tanh": {"class": "activation", "activation": "tanh", "from": ["energy_in"]},
+        "energy": {"class": "linear", "activation": None, "with_bias": False, "from": ["energy_tanh"], "n_out": 1},
+        "att_weights": {"class": "softmax_over_spatial", "from": ["energy"]},  # (B, enc-T, 1)
+        "accum_att_weights": {"class": "eval", "from": ["prev:accum_att_weights", "att_weights", "base:fertility"],
+                              "eval": "source(0) + source(1) / (2.0 * source(2))",
+                              "out_type": {"dim": 1, "shape": (None, 1)}},
+        "att": {"class": "generic_attention", "weights": "att_weights", "base": "base:encoder"},
+        "s": {"class": "rnn_cell", "unit": "standardlstm", "unit_opts": {"use_peepholes": True, "forget_bias": 0.0},
+              "initial_state": "var", "from": ["target_embed", "att"], "n_out": 10},
+        "readout_in": {"class": "linear", "from": ["prev:s", "prev:target_embed", "att"], "activation": None,
+                       "n_out": 10, "L2": l2_readout_in},
+        "readout": {"class": "reduce_out", "mode": "max", "num_pieces": 2, "from": ["readout_in"]},
+        "output_prob": {"class": "softmax", "from": ["readout"], "target": "classes", "loss": "ce"}
+      }, "target": "classes", "max_seq_len": 20},
+
+      "decision": {"class": "decide", "from": ["output"], "loss": "edit_distance", "target": "classes"}
+    }
+
+  print("Constructing train network without constraints")
+  tf.reset_default_graph()
+  extern_data = ExternData({
+    "data": {"dim": n_src_dim, "sparse": True},
+    "classes": {"dim": n_tgt_dim, "sparse": True, "available_for_inference": False}})
+  train_net = TFNetwork(
+    extern_data=extern_data, search_flag=False, train_flag=get_global_train_flag_placeholder(), config=config)
+  assert train_net.eval_flag is True
+  train_net.construct_from_dict(get_net_dict(l2_target_embed=0.0, l2_readout_in=0.0))
+  train_out_layer = train_net.layers["output"]
+  assert isinstance(train_out_layer, RecLayer)
+  assert isinstance(train_out_layer.cell, _SubnetworkRecCell)
+  assert_equal(set(train_out_layer.cell.input_layers_moved_out), {"output", "target_embed"})
+  assert_equal(set(train_out_layer.cell.output_layers_moved_out), {"output_prob", "readout_in", "readout"})
+  assert_equal(train_net.get_total_constraints(), 0)
+
+  print("Constructing train network with L2 norm on moved out input layer")
+  tf.reset_default_graph()
+  extern_data = ExternData({
+    "data": {"dim": n_src_dim, "sparse": True},
+    "classes": {"dim": n_tgt_dim, "sparse": True, "available_for_inference": False}})
+  train_net = TFNetwork(
+    extern_data=extern_data, search_flag=False, train_flag=get_global_train_flag_placeholder(), config=config)
+  assert train_net.eval_flag is True
+  train_net.construct_from_dict(get_net_dict(l2_target_embed=0.01, l2_readout_in=0.0))
+  train_out_layer = train_net.layers["output"]
+  assert isinstance(train_out_layer, RecLayer)
+  assert isinstance(train_out_layer.cell, _SubnetworkRecCell)
+  assert_equal(set(train_out_layer.cell.input_layers_moved_out), {"output", "target_embed"})
+  assert_equal(set(train_out_layer.cell.output_layers_moved_out), {"output_prob", "readout_in", "readout"})
+  assert_not_equal(train_net.get_total_constraints(), 0)
+
+  print("Constructing train network with L2 norm on moved out output layer")
+  tf.reset_default_graph()
+  extern_data = ExternData({
+    "data": {"dim": n_src_dim, "sparse": True},
+    "classes": {"dim": n_tgt_dim, "sparse": True, "available_for_inference": False}})
+  train_net = TFNetwork(
+    extern_data=extern_data, search_flag=False, train_flag=get_global_train_flag_placeholder(), config=config)
+  assert train_net.eval_flag is True
+  train_net.construct_from_dict(get_net_dict(l2_target_embed=0.0, l2_readout_in=0.01))
+  train_out_layer = train_net.layers["output"]
+  assert isinstance(train_out_layer, RecLayer)
+  assert isinstance(train_out_layer.cell, _SubnetworkRecCell)
+  assert_equal(set(train_out_layer.cell.input_layers_moved_out), {"output", "target_embed"})
+  assert_equal(set(train_out_layer.cell.output_layers_moved_out), {"output_prob", "readout_in", "readout"})
+  assert_not_equal(train_net.get_total_constraints(), 0)
+
+
 def test_rec_layer_search_select_src():
   from TFNetworkRecLayer import _SubnetworkRecCell
   n_src_dim = 5