PT engine forward task

Implement and fix #1336

Author: albertz

.github/workflows/main.yml

@@ -261,6 +261,7 @@ jobs:
         action:
           - TEST=demos RETURNN_DISABLE_TF=1
           - TEST=PTDataset
+          - TEST=torch_engine
           - TEST=torch_frontend
           - TEST=torch_internal_frontend
 

returnn/__main__.py

@@ -445,22 +445,36 @@ def execute_main_task():
             lr_control_update_scores=lr_control_update_scores,
         )
     elif task in ["forward", "hpx"]:
-        assert eval_data is not None, "no eval data provided"
-        combine_labels = config.value("combine_labels", "")
-        engine.use_search_flag = config.bool("forward_use_search", False)
-        if config.has("epoch"):
-            config.set("load_epoch", config.int("epoch", 0))
-        engine.init_network_from_config(config)
-        output_file = config.value("output_file", "dump-fwd-epoch-%i.hdf" % engine.epoch)
-        forward_batch_size = config.int("forward_batch_size", 0)
-        if not forward_batch_size:
-            raise Exception("forward_batch_size not set")
-        engine.forward_to_hdf(
-            data=eval_data,
-            output_file=output_file,
-            combine_labels=combine_labels,
-            batch_size=forward_batch_size,
-        )
+        if config.typed_value("forward_callback") or not BackendEngine.is_tensorflow_selected():
+            engine.init_network_from_config(config)
+            if config.value("forward_data", "eval") in ["train", "dev", "eval"]:
+                data = {"train": train_data, "dev": dev_data, "eval": eval_data}[config.value("forward_data", "eval")]
+                assert data, "set forward_data"
+            else:
+                data = init_dataset(config.opt_typed_value("forward_data"))
+            forward_callback = config.typed_value("forward_callback")
+            assert forward_callback, "no forward_callback specified"
+            if callable(forward_callback):
+                forward_callback = forward_callback()
+            engine.forward_with_callback(dataset=data, callback=forward_callback)
+        else:
+            assert BackendEngine.is_tensorflow_selected()
+            assert eval_data is not None, "no eval data provided"
+            combine_labels = config.value("combine_labels", "")
+            engine.use_search_flag = config.bool("forward_use_search", False)
+            if config.has("epoch"):
+                config.set("load_epoch", config.int("epoch", 0))
+            engine.init_network_from_config(config)
+            output_file = config.value("output_file", "dump-fwd-epoch-%i.hdf" % engine.epoch)
+            forward_batch_size = config.int("forward_batch_size", 0)
+            if not forward_batch_size:
+                raise Exception("forward_batch_size not set")
+            engine.forward_to_hdf(
+                data=eval_data,
+                output_file=output_file,
+                combine_labels=combine_labels,
+                batch_size=forward_batch_size,
+            )
     elif task == "search":
         engine.use_search_flag = True
         engine.use_eval_flag = config.bool("search_do_eval", True)

returnn/engine/base.py

@@ -13,6 +13,8 @@
 from returnn.log import log
 from returnn.pretrain import Pretrain
 from returnn.util import basic as util
+from returnn.forward_iface import ForwardCallbackIface
+from returnn.datasets import Dataset
 
 
 class EngineBase:
@@ -241,3 +243,11 @@ def is_first_epoch_after_pretrain(self):
         :rtype: bool
         """
         return self.pretrain and self.epoch == self.pretrain.get_train_num_epochs() + 1
+
+    def forward_with_callback(self, *, dataset: Dataset, callback: ForwardCallbackIface):
+        """
+        Iterate through the dataset, calling `forward_step` from user config,
+        collecting outputs in `rf.get_run_ctx()` via `mark_as_output` calls,
+        and then calling `callback` for each entry.
+        """
+        raise NotImplementedError

returnn/forward_iface.py

@@ -0,0 +1,39 @@
+"""
+Defines the interface for the "forward" task,
+which can be used for recognition, alignment, search, etc.
+
+https://github.com/rwth-i6/returnn/issues/1336
+"""
+
+from __future__ import annotations
+from returnn.tensor import TensorDict
+
+
+class ForwardCallbackIface:
+    """
+    Callback interface for the forward task.
+
+    Define `forward_callback` in your config to an instance or class of this.
+
+    https://github.com/rwth-i6/returnn/issues/1336
+    """
+
+    def init(self, *, model):
+        """
+        Run at the beginning.
+        """
+
+    def process_seq(self, *, seq_tag: str, outputs: TensorDict):
+        """
+        Called for each sequence, or entry in the dataset.
+        This does not have the batch dim anymore.
+        The values in `outputs` are Numpy arrays.
+
+        :param seq_tag:
+        :param outputs:
+        """
+
+    def finish(self):
+        """
+        Run at the end.
+        """

returnn/torch/engine.py

@@ -24,6 +24,8 @@
 from returnn.util import basic as util
 from returnn.util import NumbersDict
 from returnn.util.basic import NotSpecified
+from returnn.forward_iface import ForwardCallbackIface
+
 from .updater import Updater
 from .data import pipeline as data_pipeline
 from .data import returnn_dataset_wrapper
@@ -55,6 +57,11 @@ def __init__(self, config: Config):
         self._orig_model = None  # type: Optional[Union[rf.Module, torch.nn.Module]]
         self._pt_model = None  # type: Optional[torch.nn.Module]
         self._train_step_func = None  # type: Optional[Callable]
+        self._forward_step_func = self.config.typed_value("forward_step")  # type: Optional[Callable]
+        self._forward_step_expected_outputs = None  # type: Optional[TensorDict]
+        if self.config.typed_value("model_outputs") is not None:
+            self._forward_step_expected_outputs = TensorDict()
+            self._forward_step_expected_outputs.update(self.config.typed_value("model_outputs"), auto_convert=True)
         self._save_model_epoch_interval = 1
         self._updater = None  # type: Optional[Updater]
 
@@ -98,6 +105,7 @@ def init_network_from_config(self, config: Optional[Config] = None):
 
         extern_data = TensorDict()
         extern_data_dict = self.config.typed_value("extern_data")
+        assert extern_data_dict, "extern_data is not specified in config"
         extern_data.update(extern_data_dict, auto_convert=True)
         if "seq_tag" not in extern_data.data:
             batch_dim = _get_batch_dim_from_extern_data(extern_data)
@@ -194,9 +202,12 @@ def train_epoch(self):
         accumulated_losses_dict = NumbersDict()
         accumulated_inv_norm_factors_dict = NumbersDict()
         step_idx = 0
-        for data in self._train_dataloader:
+        for extern_data_raw in self._train_dataloader:
             self._updater.get_optimizer().zero_grad()
-            self._run_step(data, train_flag=True)
+            extern_data = _raw_dict_to_extern_data(
+                extern_data_raw, extern_data_template=self.extern_data, device=self._device
+            )
+            self._run_step(extern_data, train_func=True, train_flag=True)
 
             train_ctx = rf.get_run_ctx()
             total_loss = train_ctx.total_loss()
@@ -265,9 +276,12 @@ def eval_model(self):
             step_idx = 0
 
             with torch.no_grad():
-                for data in data_loader:
+                for extern_data_raw in data_loader:
+                    extern_data = _raw_dict_to_extern_data(
+                        extern_data_raw, extern_data_template=self.extern_data, device=self._device
+                    )
 
-                    self._run_step(data)
+                    self._run_step(extern_data, train_func=True)
                     train_ctx = rf.get_run_ctx()
 
                     if score_keys is None:
@@ -345,19 +359,23 @@ def _create_data_loader(self, dataset: Dataset) -> DataLoader2:
                 raise ModuleNotFoundError("Possible type error in DataLoader2 due to missing module 'dill'") from exc
             raise
 
-    def _run_step(self, extern_data_raw: Dict[str, torch.Tensor], *, train_flag: bool = False):
+    def _run_step(self, extern_data: TensorDict, *, train_flag: bool = False, train_func: bool):
         """
-        :param dict[str, torch.Tensor] extern_data_raw: model inputs for the step
+        :param extern_data: model inputs for the step
         """
-        extern_data = _raw_dict_to_extern_data(
-            extern_data_raw, extern_data_template=self.extern_data, device=self._device
-        )
-
-        rf.init_train_step_run_ctx(train_flag=train_flag)
+        if train_func:
+            assert self._train_step_func is not None
+            rf.init_train_step_run_ctx(train_flag=train_flag)
+        else:
+            assert self._forward_step_func is not None, "define forward_step in the config"
+            rf.init_forward_step_run_ctx(expected_outputs=self._forward_step_expected_outputs)
 
         with autocast(device_type=self._device, dtype=self._autocast_dtype) if self._use_autocast else nullcontext():
             sentinel_kw = {"__fwd_compatible_random_arg_%i" % int(random() * 100): None}
-            self._train_step_func(model=self._orig_model, extern_data=extern_data, **sentinel_kw)
+            if train_func:
+                self._train_step_func(model=self._orig_model, extern_data=extern_data, **sentinel_kw)
+            else:
+                self._forward_step_func(model=self._orig_model, extern_data=extern_data, **sentinel_kw)
 
     def _load_model(self):
         """
@@ -393,7 +411,7 @@ def _load_model(self):
         rf.set_random_seed(random_seed)
 
         get_model_func = self.config.typed_value("get_model")
-        assert get_model_func, "get_model not defined"
+        assert get_model_func, "get_model not defined in config"
         sentinel_kw = {"__fwd_compatible_random_arg_%i" % int(random() * 100): None}
         # Note on the `epoch` and `step` args:
         # This is the current epoch and step, i.e. the epoch and step we are about to run.
@@ -521,6 +539,42 @@ def _save_optimizer(self):
             if os.path.isfile(filename):
                 os.unlink(filename)
 
+    def forward_with_callback(self, *, dataset: Dataset, callback: ForwardCallbackIface):
+        """forward"""
+        assert isinstance(dataset, Dataset)
+        assert isinstance(callback, ForwardCallbackIface)
+
+        self._pt_model.eval()
+
+        data_loader = self._create_data_loader(dataset)
+        batch_dim = _get_batch_dim_from_extern_data(self.extern_data)
+
+        with torch.no_grad():
+            callback.init(model=self._orig_model)
+
+            for extern_data_raw in data_loader:
+                extern_data = _raw_dict_to_extern_data(
+                    extern_data_raw, extern_data_template=self.extern_data, device=self._device
+                )
+                self._run_step(extern_data, train_func=False)
+                ctx = rf.get_run_ctx()
+                ctx.check_outputs_complete()
+
+                model_outputs = ctx.outputs
+                model_outputs_per_batch_template = TensorDict(
+                    {k: v.copy_template_excluding_axis(0) for k, v in model_outputs.data.items()}
+                )
+                for batch_idx in range(batch_dim.get_dim_value()):
+                    seq_tag = extern_data["seq_tag"].raw_tensor[batch_idx].item()
+                    model_outputs_per_batch = TensorDict(
+                        {k: v.copy() for k, v in model_outputs_per_batch_template.data.items()}
+                    )
+                    for k, v in model_outputs.data.items():
+                        model_outputs_per_batch[k].raw_tensor = v.raw_tensor[batch_idx]
+                    callback.process_seq(seq_tag=seq_tag, outputs=model_outputs_per_batch)
+
+            callback.finish()
+
 
 def _to_raw(n: Union[int, float, Tensor]):
     if isinstance(n, (int, float)):

tests/test_torch_engine.py

@@ -0,0 +1,53 @@
+"""
+Tests for PyTorch engine.
+"""
+
+import _setup_test_env  # noqa
+import torch
+from returnn.config import Config, global_config_ctx
+from returnn.tensor import TensorDict, Tensor
+from returnn.torch.engine import Engine
+import returnn.frontend as rf
+from returnn.forward_iface import ForwardCallbackIface
+from returnn.datasets import init_dataset
+
+
+def test_torch_engine():
+    def _get_model(**_kwargs):
+        return torch.nn.Module()
+
+    def _forward_step(*, extern_data: TensorDict, **_kwargs):
+        rf.get_run_ctx().mark_as_default_output(extern_data["data"])
+
+    class _ForwardCallback(ForwardCallbackIface):
+        def __init__(self):
+            self.num_seqs = 0
+            self.init_called = False
+            self.finish_called = False
+
+        def init(self, *, model):
+            assert isinstance(model, torch.nn.Module)
+            assert self.num_seqs == 0
+            self.init_called = True
+
+        def process_seq(self, *, seq_tag: str, outputs: TensorDict):
+            out = outputs["output"]
+            assert isinstance(out, Tensor)
+            assert out.batch_ndim == 2 and out.batch_shape[-1] == 9
+            self.num_seqs += 1
+
+        def finish(self):
+            self.finish_called = True
+
+    config = Config(
+        dict(task="forward", extern_data={"data": {"dim": 9}}, get_model=_get_model, forward_step=_forward_step)
+    )
+    dataset = init_dataset({"class": "Task12AXDataset", "num_seqs": 100, "name": "dev", "fixed_random_seed": 1})
+    callback = _ForwardCallback()
+
+    with global_config_ctx(config):
+        engine = Engine(config=config)
+        engine.init_network_from_config()
+        engine.forward_with_callback(callback=callback, dataset=dataset)
+        assert callback.num_seqs == 100
+        assert callback.init_called and callback.finish_called