Repository: rwth-i6/returnn

Issue:
Choice layer is not capable of simply forwarding input
I have a test case for my scenario here:
https://github.com/JackTemaki/returnn/commit/9468b301289d2df6d2b5c068babc859f0c3235e8

So if there is a network without beam-search, that provides just the feedback of any target via choice layer, the network can not be used with `task="forward"`. 

RETURNN will move layers out of the loop which are supposed to be in the loop:
```
v3: DEBUG: Rec layer 'output' (search False, train False) sub net:
v3: DEBUG:   Input layers moved out of loop: (#: 2)
v3: DEBUG:     choice
v3: DEBUG:     feedback
v3: DEBUG:   Output layers moved out of loop: (#: 1)
v3: DEBUG:     output
v3: DEBUG:   Layers in loop: (#: 6)
v3: DEBUG:     end
v3: DEBUG:     end_compare
v3: DEBUG:     stop_token_sigmoid
v3: DEBUG:     s
v3: DEBUG:     c
v3: DEBUG:     c_in
v3: DEBUG:   Unused layers: (#: 0)
v3: DEBUG:     None
```
And then fail the construction (the actual error is not relevant here)

The Problem is, that there is no way to tell a choice layer that it should simply copy the input during decoding. An alternative is to enable the search_flag, which will result in creating a beam which is wrong here, and will also lead to another crash.

There is also a comment in the code suggesting that this is not intended to work at all [here](https://github.com/rwth-i6/returnn/blob/master/returnn/tf/layers/rec.py#L4551):
```python
    else:  # no search, and no scheduled-sampling
      if not self.network.search_flag:
        assert len(self.sources) == 0  # will be filtered out in transform_config_dict
      # Note: If you want to do forwarding, without having the reference,
      # that wont work. You must do search in that case.
```

So what is the correct approach to feed the ground truth or do scheduled sampling in training, but simply copying the input during forwarding?




Discussion:
I would guess there is a way to make it work with search and beam size 1, but yes, this should not be needed.
Forwarding is simply not implemented in the layer, probably because in case of `input_type="prob"` it is not clear what forwarding means (maybe greedy search?). Also ChoiceLayer is a misnomer for regression, I guess it is not a separate layer because of the scheduled sampling implementation?

I think we have to put the handling of regression in the search code instead into a `not self.train_flag and input_type == "regression"` block and then inside only handle the search choices `if self.search_flag`.
I'm already working on it.

The intended behavior of the choice layer is to return the generated/predicted output with search flag enabled, and otherwise the ground truth.

Resolving patch:
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
@@ -2115,6 +2115,8 @@ def get_loop_loss():
       if self.input_layers_moved_out:
         with tf.name_scope("input_layers_moved_out"):
           self._construct_input_layers_moved_out()
+          if fixed_seq_len is None and rec_layer.output.size_placeholder:  # might have set it by now
+            fixed_seq_len = rec_layer.output.size_placeholder[0]
           for layer_name in self.input_layers_moved_out:
             # Create only Tensor arrays for those which we use inside the loop.
             if not self._input_layer_used_inside_loop(layer_name):
@@ -2123,11 +2125,6 @@ def get_loop_loss():
             assert isinstance(layer, LayerBase)
             if layer_name == "output":
               assert layer.output.have_time_axis()
-              # If we don't know our own size yet, we can overtake it from this layer.
-              if not rec_layer.output.size_placeholder:
-                rec_layer.output.size_placeholder = {0: layer.output.get_sequence_lengths()}
-              if fixed_seq_len is None:
-                fixed_seq_len = layer.output.get_sequence_lengths()
               assert rec_layer.output.is_same_time_dim(layer.output)
             # Only unroll if that is the same time dim.
             if not layer.output.mark_same_time(rec_layer.output):
@@ -2269,7 +2266,8 @@ def maybe_transform(layer):
 
         if seq_len_info is not None:
           assert self.net.layers["end"].output.shape == (), "end layer %r unexpected shape" % self.net.layers["end"]
-          choices = self.net.layers["end"].get_search_choices()
+          end_layer = maybe_transform(self.net.layers["end"])
+          choices = end_layer.get_search_choices()
           if choices:
             from .basic import SelectSearchSourcesLayer
             cur_end_layer = choices.translate_to_this_search_beam(prev_end_layer)
@@ -2277,7 +2275,7 @@ def maybe_transform(layer):
               "unexpected search choices: cur end %r, prev end %r" % (choices, prev_end_layer.get_search_choices()))
             assert cur_end_layer.search_choices_seq, (
               "unexpected search choices: cur end %r (via %r), prev end %r (via %r)" % (
-                choices, self.net.layers["end"], prev_end_layer.get_search_choices(), prev_end_layer))
+                choices, end_layer, prev_end_layer.get_search_choices(), prev_end_layer))
             assert cur_end_layer.output.shape == (), "end layer %r unexpected shape" % cur_end_layer
             with tf.name_scope("end_flag"):
               end_flag = cur_end_layer.output.placeholder
@@ -2597,7 +2595,12 @@ def get_choice_seq(choice_base):
       end_layer = None
 
     if end_layer:
+      # seq_len is determined from the end-layer. We need to translate it to the right beam.
       end_layer_choice = self.net.get_search_choices(src=end_layer)
+      assert end_layer_choice and end_layer_choice.search_choices
+      if end_layer_choice.name.startswith("prev:"):
+        # Logic from maybe_transform. It would be translated to the current beam.
+        end_layer_choice = self.net.layers[end_layer_choice.name[len("prev:"):]]
       assert end_layer_choice in choice_seq_in_frame, (
         "End layer must not have a beam independent from output layer '{}'.".format(layer_name))
 
@@ -2985,6 +2988,19 @@ def get_layer(name):
       for layer_name in self.input_layers_moved_out:
         get_layer(layer_name)
 
+    # We might have figured out the real output seq length (and dim tag) by now.
+    if not self.parent_rec_layer.output.size_placeholder and "output" in self.input_layers_moved_out:
+      output_layer = self.input_layers_net.layers["output"]
+      assert output_layer.output.have_time_axis()
+      self.parent_rec_layer.output.size_placeholder = {0: output_layer.output.get_sequence_lengths()}
+    # This might be set e.g. by ChoiceLayer, or losses.
+    if not self.parent_rec_layer.output.size_placeholder and self.input_layers_net.used_data_keys:
+      for data_key in sorted(self.input_layers_net.used_data_keys):
+        data = self.input_layers_net.extern_data.data[data_key]
+        if data.have_time_axis():
+          self.parent_rec_layer.output.size_placeholder = {0: data.get_sequence_lengths()}
+          break
+
   def _construct_output_layers_moved_out(self, loop_accumulated, seq_len, extra_output_layers, final_net_vars):
     """
     See self._move_outside_loop().
@@ -3645,6 +3661,7 @@ def __init__(self, n_out, unit, unit_opts=None,
         if isinstance(self.cell, BaseRNNCell):
           x = self.cell.get_input_transformed(x)
         assert not self.input_data or self.input_data.time_dim_axis is None, (
+          self, self.input_data,
           "A recurrent layer is not allowed to have input data with a remaining time axis.\n"
           "A possible reason for this error is that the 'target' of the rec layer does not\n"
           "match the targets of the sub-layers")


Test patch:
diff --git a/tests/test_TFEngine.py b/tests/test_TFEngine.py
--- a/tests/test_TFEngine.py
+++ b/tests/test_TFEngine.py
@@ -3534,6 +3534,84 @@ def test_non_available_data_construction():
   engine.finalize()
 
 
+def test_regression_choice():
+  net_dict = {
+    "encoder": {"class": "linear", "activation": "tanh", "from": ["data:classes"], "n_out": 20},
+    "enc_ctx": {"class": "linear", "activation": "tanh", "from": ["encoder"], "n_out": 20},
+
+    "output": {"class": "rec", "from": [], "max_seq_len": 10, "unit": {
+      "s": {"class": "rnn_cell", "unit": "LSTMBlock", "from": ["prev:c", "feedback"], "n_out": 20},
+      "c_in": {"class": "linear", "activation": "tanh", "from": ["s"], "n_out": 20},
+      "c": {"class": "dot_attention", "from": ["c_in"], "base": "base:encoder", "base_ctx": "base:enc_ctx",
+            "n_out": 20},
+      "output": {"class": "linear", "from": ["s", "c"], "target": "data", 'n_out': 2, 'activation': None,
+                 "loss": "mse"},
+      "choice": {'class': 'choice', 'beam_size': 1, 'input_type': 'regression', 'from': 'output',
+                 'target': 'data', 'n_out': 2},
+      "feedback": {'class': 'linear', 'from': 'prev:choice', 'n_out': 5, 'activation': 'tanh'},
+      'end_compare': {'class': 'compare', 'from': ['stop_token_sigmoid'], 'kind': 'greater', 'value': 0.5},
+      'end': {'class': 'squeeze', 'from': ['end_compare'], 'axis': 'F'},
+      'stop_token_sigmoid': {'activation': 'sigmoid', 'class': 'linear', 'from': ['s'], 'n_out': 1}
+    }, "target": "data"}
+
+  }
+
+  from returnn.datasets.generating import DummyDataset, StaticDataset
+  seq_len = 5
+  n_data_dim = 2
+  n_classes_dim = 3
+  dataset = DummyDataset(input_dim=n_data_dim, output_dim=n_classes_dim, num_seqs=2, seq_len=seq_len)
+  dataset.init_seq_order(epoch=1)
+  search_data = StaticDataset(
+    data=[{'classes': numpy.ones((seq_len,), dtype="int32")}], output_dim={'classes': (n_classes_dim, 1)})
+  search_data.init_seq_order(epoch=1)
+
+  import tempfile
+  output_file = tempfile.mktemp(suffix=".hdf", prefix="nose-tf-forward")
+
+  config = Config()
+  config.update({
+    "model": "%s/model" % _get_tmp_dir(),
+    "batch_size": 100,
+    "max_seqs": 2,
+    "extern_data": {
+      'data': {'dim': n_data_dim, 'shape': (None, n_data_dim), 'available_for_inference': False},
+      'classes': {'dim': n_classes_dim, 'shape': (None,), 'sparse': True, 'available_for_inference': True}},
+    "network": net_dict,
+    "start_epoch": 1,
+    "num_epochs": 2,
+    "learning_rate": 0.01,
+    "adam": True,
+    "debug_print_layer_output_template": True,
+    "debug_print_layer_output_shape": True,
+  })
+
+  _cleanup_old_models(config)
+  engine = Engine(config=config)
+  engine.init_train_from_config(config=config, train_data=dataset, dev_data=dataset, eval_data=None)
+  print("Extern data:")
+  pprint(engine.network.extern_data.data)
+  print("Used data keys:")
+  pprint(engine.network.used_data_keys)
+  engine.train()
+
+  print("Search...")
+  engine.use_search_flag = True
+  engine.use_eval_flag = False
+  engine.use_dynamic_train_flag = False
+  engine.init_network_from_config(config)
+  engine.search(dataset=search_data, do_eval=False)
+
+  print("Forward...")
+  engine.use_search_flag = False
+  engine.use_eval_flag = False
+  engine.use_dynamic_train_flag = False
+  engine.init_network_from_config(config)
+  engine.forward_to_hdf(data=dataset, output_file=output_file, batch_size=1)
+
+  engine.finalize()
+
+
 if __name__ == "__main__":
   try:
     better_exchook.install()