PyTorch search interface

[albertz]

This is for doing beam search inside RETURNN, e.g. for end-to-end models, e.g. AED.
This is when using the PyTorch engine (#1120).

Currently when the user uses task = "search", our main would call engine.search(...). This function is not implemented for the PT engine.

Do we want to implement it?

In #1120, currently we discussed to have a very generic forward_step function, which can have any number of outputs. But we don't really have discussed how they would be used. Dumped to HDF? Dumped to a Python file similar as what we get via TF engine.search? How would that be configured?

The TF engine.search is more specific. Either the search output is behind DecideLayer, so no beam anymore, then it will create a Python file with a dict seq_tag -> hyp. Or it is with beam, then it is a Python file with a dict seq_tag -> list of tuple hyp and score.

Do we want to have sth as specific as that? So the user defines a search_step or so in the config?

Or how exactly would we define in the config what would happen with the outputs from forward_step? All dumped to file? But this should be flexible enough to allow for a format exactly like the TF format. I.e. Python file with a dict seq_tag -> list of tuple hyp and score.

[JackTemaki]

For me it would be fine to have nothing implemented, I prefer to take care of any processing and dumping myself in the forward_step, I would not even require a "search" task.

[albertz]

Currently when the user puts task = "forward", RETURNN calls engine.forward_to_hdf(...). So as the name says, it is already somewhat specific to HDF. And it is expecting a single output.

Maybe to have it more distinct, we just introduce a new more generic task = "forward_generic" or so.

But then still the question remains, how would the output(s) from forward_step ends up in a file. How would the user specify this in the config? Ideally this should be generic enough that it is possible to generate a file with the same format as what we get via TF engine.search. I.e. Python file with a dict seq_tag -> list of tuple hyp and score.

[albertz]

Maybe the user can define some forward_callback(seq_tag, outputs) or so which is executed after the forward_step, where we pass the final outputs. Maybe this is called already per each sequence, no batch dim anymore. Then the user could implement the file saving by hand in there. Maybe an initial forward_initial_callback to open the files. Or maybe directly the user passes in some object derived from interface:

class ForwardCallbackIface:
  def init(self): pass
  def process_seq(self, seq_tag, outputs): pass
  def finish(self): pass

Or so. Then the user specifies forward_callback = MyForwardCallback in the config.

If forward_callback is callable (it is here because of class), it will call it (then it gets MyForwardCallback(), i.e. an instance of MyForwardCallback in this example).

MyForwardCallback could look like this:

class MyForwardCallback(ForwardCallbackIface):
  def __init__(self):
    self.out_file = None
  def init(self):
    self.out_file = gzip.open("recog.out.py.gz", "wt")
    self.out_file.write("{\n")
  def process_seq(self, seq_tag, outputs):
    self.out_file.write(repr(seq_tag) + ": [%s],\n" % ...)
  def finish(self):
    self.out_file.write("}\n")
    self.out_file.close()

[albertz]

While this is very generic, maybe it's too complicated for the user? Is there some easier way for the user? (Which should still be straightforward to use.)

Maybe for the task = "search" case, we still want to have this more specific variant, that it is simpler, and expects exactly hyps, scores as output, in a well defined way, and then follows the existing TF output formats? (For now maybe only the Python format case, not sure if we want the others.)

[albertz]

Questions:

• Is the suggested interface ForwardStepCallbackIntf ok? Or how should it look like instead? -> ForwardCallbackIface, step -> process_seq
• Is the option forward_step_callback to set this ok? -> forward_callback
• Is task = "forward_generic" good for it? Or how should we call it instead? -> task = "forward"
• We still would use the forward_step function? Or how should we call it instead? -> forward_step

[albertz]

As there is no other suggestion, and we already waited here for almost a week, I'm going to implement those now. If you want to have anything different, please respond as soon as possible.