Repository: iterait/emloop

Issue:
Specify the name of the training stream in the config
Migrated from [https://github.com/Cognexa/cxflow/issues/190](https://github.com/Cognexa/cxflow/issues/190).

It would be pretty nice to have a single dataset with let's say two trainign streams (generative and regular one). One would be able to pre-train the model on the generated data and then (in the separate training) finetune it on the real data.

Implementing this would make cxflow be more consistent with other streams that are easily configurable.

Config suggestion:
```
main_loop:
  train_stream: train  # or generative_train
```

Discussion:

Patch:
diff --git a/emloop/constants.py b/emloop/constants.py
index 0249a86..5ebf78b 100644
--- a/emloop/constants.py
+++ b/emloop/constants.py
@@ -33,9 +33,9 @@ EL_NA_STR = 'N/A'
 EL_BUFFER_SLEEP = 0.02
 """The duration for which the buffer sleeps before it starts to process the next batch."""
 
-EL_TRAIN_STREAM = 'train'
+EL_DEFAULT_TRAIN_STREAM = 'train'
 """The stream to be used for training."""
 
 __all__ = ['EL_LOG_FORMAT', 'EL_LOG_DATE_FORMAT', 'EL_FULL_DATE_FORMAT', 'EL_HOOKS_MODULE', 'EL_CONFIG_FILE',
-           'EL_LOG_FILE', 'EL_TRACE_FILE', 'EL_TRAIN_STREAM', 'EL_PREDICT_STREAM', 'EL_DEFAULT_LOG_DIR',
+           'EL_LOG_FILE', 'EL_TRACE_FILE', 'EL_DEFAULT_TRAIN_STREAM', 'EL_PREDICT_STREAM', 'EL_DEFAULT_LOG_DIR',
            'EL_NA_STR', 'EL_BUFFER_SLEEP']
diff --git a/emloop/hooks/abstract_hook.py b/emloop/hooks/abstract_hook.py
index 039f3cd..7ecb928 100644
--- a/emloop/hooks/abstract_hook.py
+++ b/emloop/hooks/abstract_hook.py
@@ -81,7 +81,7 @@ class AbstractHook:
         """
         pass
 
-    def after_epoch_profile(self, epoch_id: int, profile: TimeProfile, extra_streams: Iterable[str]) -> None:
+    def after_epoch_profile(self, epoch_id: int, profile: TimeProfile, train_stream_name: str, extra_streams: Iterable[str]) -> None:
         """
         After epoch profile event.
 
diff --git a/emloop/hooks/log_profile.py b/emloop/hooks/log_profile.py
index 028bb61..dd77b55 100644
--- a/emloop/hooks/log_profile.py
+++ b/emloop/hooks/log_profile.py
@@ -24,13 +24,13 @@ class LogProfile(AbstractHook):
 
     """
 
-    def after_epoch_profile(self, epoch_id, profile: TimeProfile, extra_streams: Iterable[str]) -> None:
+    def after_epoch_profile(self, epoch_id, profile: TimeProfile, train_stream_name: str, extra_streams: Iterable[str]) -> None:
         """
         Summarize and log the given epoch profile.
 
         The profile is expected to contain at least:
             - ``read_data_train``, ``eval_batch_train`` and ``after_batch_hooks_train`` entries produced by the train
-              stream
+              stream (if train stream name is `train`)
             - ``after_epoch_hooks`` entry
 
         :param profile: epoch timings profile
@@ -39,13 +39,13 @@ class LogProfile(AbstractHook):
 
         read_data_total = 0
         eval_total = 0
-        train_total = sum(profile.get('eval_batch_train', []))
+        train_total = sum(profile.get('eval_batch_{}'.format(train_stream_name), []))
         hooks_total = sum(profile.get('after_epoch_hooks', []))
 
-        for stream_name in chain(extra_streams, ['train']):
+        for stream_name in chain(extra_streams, [train_stream_name]):
             read_data_total += sum(profile.get('read_batch_' + stream_name, []))
             hooks_total += sum(profile.get('after_batch_hooks_' + stream_name, []))
-            if stream_name != 'train':
+            if stream_name != train_stream_name:
                 eval_total += sum(profile.get('eval_batch_' + stream_name, []))
 
         logging.info('\tT read data:\t%f', read_data_total)
diff --git a/emloop/hooks/stop_after.py b/emloop/hooks/stop_after.py
index c0d821f..d2c2348 100644
--- a/emloop/hooks/stop_after.py
+++ b/emloop/hooks/stop_after.py
@@ -6,7 +6,7 @@ from datetime import datetime
 from typing import Optional
 
 from . import AbstractHook, TrainingTerminated
-from ..constants import EL_TRAIN_STREAM
+from ..constants import EL_DEFAULT_TRAIN_STREAM
 from ..types import EpochData, Batch
 
 
@@ -32,7 +32,7 @@ class StopAfter(AbstractHook):
     """
 
     def __init__(self, epochs: Optional[int]=None, iterations: Optional[int]=None, minutes: Optional[float]=None,
-                 **kwargs):
+                 train_stream_name: str=EL_DEFAULT_TRAIN_STREAM, **kwargs):
         """
         Create new StopAfter hook.
 
@@ -67,6 +67,7 @@ class StopAfter(AbstractHook):
         self._minutes = minutes
         self._iters_done = 0
         self._training_start = None
+        self._train_stream_name = train_stream_name
 
     def _check_train_time(self) -> None:
         """
@@ -91,7 +92,7 @@ class StopAfter(AbstractHook):
         :raise TrainingTerminated: if the number of iterations reaches ``self._iters``
         """
         self._check_train_time()
-        if self._iters is not None and stream_name == EL_TRAIN_STREAM:
+        if self._iters is not None and stream_name == self._train_stream_name:
             self._iters_done += 1
             if self._iters_done >= self._iters:
                 raise TrainingTerminated('Training terminated after iteration {}'.format(self._iters_done))
diff --git a/emloop/main_loop.py b/emloop/main_loop.py
index ba8fcce..a0d906b 100644
--- a/emloop/main_loop.py
+++ b/emloop/main_loop.py
@@ -14,7 +14,7 @@ from .hooks.abstract_hook import AbstractHook, TrainingTerminated
 from .utils import Timer, TrainingTrace, TrainingTraceKeys
 from .utils.misc import CaughtInterrupts
 from .datasets.stream_wrapper import StreamWrapper
-from .constants import EL_TRAIN_STREAM, EL_PREDICT_STREAM
+from .constants import EL_DEFAULT_TRAIN_STREAM, EL_PREDICT_STREAM
 from .types import EpochData
 
 
@@ -29,6 +29,7 @@ class MainLoop(CaughtInterrupts):   # pylint: disable=too-many-instance-attribut
     def __init__(self,   # pylint: disable=too-many-arguments
                  model: AbstractModel, dataset: AbstractDataset,
                  hooks: Iterable[AbstractHook]=(),
+                 train_stream_name: str=EL_DEFAULT_TRAIN_STREAM,
                  extra_streams: List[str]=(),  # pylint: disable=invalid-sequence-index
                  buffer: int=0,
                  on_empty_batch: str='error',
@@ -41,6 +42,7 @@ class MainLoop(CaughtInterrupts):   # pylint: disable=too-many-instance-attribut
         :param model: trained model
         :param dataset: loaded dataset
         :param hooks: training hooks
+        :param train_stream_name: name of the training stream
         :param extra_streams: additional stream names to be evaluated between epochs
         :param buffer: size of the batch buffer, 0 means no buffer
         :param on_empty_batch: action to take when batch is empty; one of :py:attr:`MainLoop.EMPTY_ACTIONS`
@@ -68,6 +70,7 @@ class MainLoop(CaughtInterrupts):   # pylint: disable=too-many-instance-attribut
         self._fixed_epoch_size = fixed_epoch_size
         self._extra_sources_warned = False
         self._epoch_profile = {}
+        self._train_stream_name = train_stream_name
         self._extra_streams = list(extra_streams)
         self._skip_zeroth_epoch = skip_zeroth_epoch
         self._streams = {}
@@ -93,7 +96,7 @@ class MainLoop(CaughtInterrupts):   # pylint: disable=too-many-instance-attribut
     def _create_epoch_data(self, streams: Optional[Iterable[str]]=None) -> EpochData:
         """Create empty epoch data double dict."""
         if streams is None:
-            streams = [EL_TRAIN_STREAM] + self._extra_streams
+            streams = [self._train_stream_name] + self._extra_streams
         return OrderedDict([(stream_name, OrderedDict()) for stream_name in streams])
 
     def _check_sources(self, batch: Dict[str, object]) -> None:
@@ -206,7 +209,7 @@ class MainLoop(CaughtInterrupts):   # pylint: disable=too-many-instance-attribut
             try:
                 stream_fn = getattr(self._dataset, stream_fn_name)
                 stream_epoch_limit = -1
-                if self._fixed_epoch_size is not None and stream_name == EL_TRAIN_STREAM:
+                if self._fixed_epoch_size is not None and stream_name == self._train_stream_name:
                     stream_epoch_limit = self._fixed_epoch_size
                 self._streams[stream_name] = StreamWrapper(stream_fn, buffer_size=self._buffer,
                                                            epoch_size=stream_epoch_limit, name=stream_name,
@@ -264,7 +267,7 @@ class MainLoop(CaughtInterrupts):   # pylint: disable=too-many-instance-attribut
             - :py:meth:`emloop.hooks.AbstractHook.after_epoch`
             - :py:meth:`emloop.hooks.AbstractHook.after_epoch_profile`
         """
-        for stream_name in [EL_TRAIN_STREAM] + self._extra_streams:
+        for stream_name in [self._train_stream_name] + self._extra_streams:
             self.get_stream(stream_name)
 
         def training():
@@ -274,7 +277,7 @@ class MainLoop(CaughtInterrupts):   # pylint: disable=too-many-instance-attribut
             # Zeroth epoch: after_epoch
             if not self._skip_zeroth_epoch:
                 logging.info('Evaluating 0th epoch')
-                self._run_zeroth_epoch([EL_TRAIN_STREAM] + self._extra_streams)
+                self._run_zeroth_epoch([self._train_stream_name] + self._extra_streams)
                 logging.info('0th epoch done\n\n')
 
             # Training loop: after_epoch, after_epoch_profile
@@ -284,7 +287,7 @@ class MainLoop(CaughtInterrupts):   # pylint: disable=too-many-instance-attribut
                 self._epoch_profile.clear()
                 epoch_data = self._create_epoch_data()
 
-                with self.get_stream(EL_TRAIN_STREAM) as stream:
+                with self.get_stream(self._train_stream_name) as stream:
                     self.train_by_stream(stream)
 
                 for stream_name in self._extra_streams:
@@ -297,6 +300,7 @@ class MainLoop(CaughtInterrupts):   # pylint: disable=too-many-instance-attribut
 
                 for hook in self._hooks:
                     hook.after_epoch_profile(epoch_id=epoch_id, profile=self._epoch_profile,
+                                             train_stream_name=self._train_stream_name,
                                              extra_streams=self._extra_streams)
                 self._epochs_done = epoch_id
                 if trace is not None:


Test patch:
diff --git a/emloop/tests/hooks/log_profile_test.py b/emloop/tests/hooks/log_profile_test.py
index 0d79e1f..e767f4b 100644
--- a/emloop/tests/hooks/log_profile_test.py
+++ b/emloop/tests/hooks/log_profile_test.py
@@ -2,6 +2,7 @@
 Module with profile hook test case (see :py:class:`emloop.hooks.LogProfile`).
 """
 import logging
+from emloop.constants import EL_DEFAULT_TRAIN_STREAM
 from emloop.hooks import LogProfile
 
 
@@ -45,7 +46,7 @@ class TestLogProfile:
         """Test KeyError raised on missing profile entries."""
         caplog.set_level(logging.INFO)
 
-        self._hook.after_epoch_profile(0, {}, [])
+        self._hook.after_epoch_profile(0, {}, EL_DEFAULT_TRAIN_STREAM, [])
         assert caplog.record_tuples == [
             ('root', logging.INFO, '\tT read data:\t0.000000'),
             ('root', logging.INFO, '\tT train:\t0.000000'),
@@ -54,7 +55,7 @@ class TestLogProfile:
         ]
 
         caplog.clear()
-        self._hook.after_epoch_profile(0, {'some_contents': 1}, [])
+        self._hook.after_epoch_profile(0, {'some_contents': 1}, EL_DEFAULT_TRAIN_STREAM, [])
         assert caplog.record_tuples == [
             ('root', logging.INFO, '\tT read data:\t0.000000'),
             ('root', logging.INFO, '\tT train:\t0.000000'),
@@ -66,7 +67,7 @@ class TestLogProfile:
         """Test profile handling with only train stream."""
         caplog.set_level(logging.INFO)
 
-        self._hook.after_epoch_profile(1, _TRAIN_ONLY_PROFILE, [])
+        self._hook.after_epoch_profile(1, _TRAIN_ONLY_PROFILE, EL_DEFAULT_TRAIN_STREAM, [])
         assert caplog.record_tuples == [
             ('root', logging.INFO, '\tT read data:\t6.120000'),
             ('root', logging.INFO, '\tT train:\t21.900000'),
@@ -78,7 +79,7 @@ class TestLogProfile:
         """Test extra streams handling."""
         caplog.set_level(logging.INFO)
 
-        self._hook.after_epoch_profile(0, _TRAIN_ONLY_PROFILE, ['valid'])
+        self._hook.after_epoch_profile(0, _TRAIN_ONLY_PROFILE, EL_DEFAULT_TRAIN_STREAM, ['valid'])
         assert caplog.record_tuples == [
             ('root', logging.INFO, '\tT read data:\t6.120000'),
             ('root', logging.INFO, '\tT train:\t21.900000'),
@@ -88,7 +89,7 @@ class TestLogProfile:
 
         # test additional entries being ignored if the extra stream was not specified
         caplog.clear()
-        self._hook.after_epoch_profile(1, _TRAIN_AND_VALID_PROFILE, [])
+        self._hook.after_epoch_profile(1, _TRAIN_AND_VALID_PROFILE, EL_DEFAULT_TRAIN_STREAM, [])
         assert caplog.record_tuples == [
             ('root', logging.INFO, '\tT read data:\t6.001000'),
             ('root', logging.INFO, '\tT train:\t21.540000'),
@@ -98,7 +99,7 @@ class TestLogProfile:
 
         # test one additional stream
         caplog.clear()
-        self._hook.after_epoch_profile(1, _TRAIN_AND_VALID_PROFILE, ['valid'])
+        self._hook.after_epoch_profile(1, _TRAIN_AND_VALID_PROFILE, EL_DEFAULT_TRAIN_STREAM, ['valid'])
         assert caplog.record_tuples == [
             ('root', logging.INFO, '\tT read data:\t20.001000'),
             ('root', logging.INFO, '\tT train:\t21.540000'),
@@ -108,7 +109,7 @@ class TestLogProfile:
 
         # test two additional streams
         caplog.clear()
-        self._hook.after_epoch_profile(1, _TRAIN_TEST_AND_VALID_PROFILE, ['valid', 'test'])
+        self._hook.after_epoch_profile(1, _TRAIN_TEST_AND_VALID_PROFILE, EL_DEFAULT_TRAIN_STREAM, ['valid', 'test'])
         assert caplog.record_tuples == [
             ('root', logging.INFO, '\tT read data:\t23.322000'),
             ('root', logging.INFO, '\tT train:\t21.540000'),
diff --git a/emloop/tests/hooks/stop_after_test.py b/emloop/tests/hooks/stop_after_test.py
index 22c3daa..80f806f 100644
--- a/emloop/tests/hooks/stop_after_test.py
+++ b/emloop/tests/hooks/stop_after_test.py
@@ -6,11 +6,11 @@ import time
 from emloop.main_loop import MainLoop
 from emloop.hooks.stop_after import StopAfter
 from emloop.hooks.abstract_hook import TrainingTerminated
-from emloop.constants import EL_TRAIN_STREAM
+from emloop.constants import EL_DEFAULT_TRAIN_STREAM
 import pytest
 
 NOTRAIN_STREAM_NAME = 'valid'
-assert EL_TRAIN_STREAM is not NOTRAIN_STREAM_NAME
+assert EL_DEFAULT_TRAIN_STREAM != NOTRAIN_STREAM_NAME
 
 
 def test_no_conditions_raise():
@@ -41,13 +41,13 @@ def test_stop_after_iters():
         hook.after_batch(stream_name=NOTRAIN_STREAM_NAME, batch_data=None)
 
     for i in range(9):
-        hook.after_batch(stream_name=EL_TRAIN_STREAM, batch_data=None)
+        hook.after_batch(stream_name=EL_DEFAULT_TRAIN_STREAM, batch_data=None)
 
     hook.after_batch(stream_name=NOTRAIN_STREAM_NAME, batch_data=None)
 
     # Test hook does terminate the training correctly
     with pytest.raises(TrainingTerminated):
-        hook.after_batch(stream_name=EL_TRAIN_STREAM, batch_data=None)
+        hook.after_batch(stream_name=EL_DEFAULT_TRAIN_STREAM, batch_data=None)
 
 
 def test_stop_after_minutes():
diff --git a/emloop/tests/main_loop_test.py b/emloop/tests/main_loop_test.py
index 0e30fdc..32a3210 100644
--- a/emloop/tests/main_loop_test.py
+++ b/emloop/tests/main_loop_test.py
@@ -10,7 +10,7 @@ import logging
 import numpy as np
 
 import emloop as el
-from emloop.constants import EL_BUFFER_SLEEP, EL_PREDICT_STREAM
+from emloop.constants import EL_BUFFER_SLEEP, EL_PREDICT_STREAM, EL_DEFAULT_TRAIN_STREAM
 from emloop.datasets import StreamWrapper
 from emloop.hooks import StopAfter
 from emloop.types import EpochData, Batch, Stream, TimeProfile
@@ -124,6 +124,17 @@ class AllEmptyBatchDataset(SimpleDataset):
             yield {'input': [], 'target': []}
 
 
+class GeneratedStreamDataset(SimpleDataset):
+    """SimpleDataset with generated stream instead of train stream."""
+
+    def generated_stream(self) -> Stream:
+        for _ in range(self.iters):
+            yield {'input': np.ones(self.shape), 'target': np.zeros(self.shape)}
+
+    def train_stream(self) -> Stream:
+        assert False
+
+
 class EventRecordingHook(el.AbstractHook):
     """EventRecordingHook records all the events and store their count and order."""
 
@@ -148,7 +159,8 @@ class EventRecordingHook(el.AbstractHook):
         self.after_epoch_events.append(self._event_id)
         self._event_id += 1
 
-    def after_epoch_profile(self, epoch_id: int, profile: TimeProfile, extra_streams: Iterable[str]) -> None:
+    def after_epoch_profile(self, epoch_id: int, profile: TimeProfile, train_stream_name: str,
+                            extra_streams: Iterable[str]) -> None:
         self.after_epoch_profile_events.append(self._event_id)
         self._event_id += 1
 
@@ -189,7 +201,8 @@ class SaveProfileHook(el.AbstractHook):
         super().__init__()
         self.profile = None
 
-    def after_epoch_profile(self, epoch_id: int, profile: TimeProfile, extra_streams: Iterable[str]) -> None:
+    def after_epoch_profile(self, epoch_id: int, profile: TimeProfile, train_stream_name: str,
+                            extra_streams: Iterable[str]) -> None:
         """Save the profile to self.profile."""
         self.profile = profile
 
@@ -280,7 +293,7 @@ class RecordingModel(TrainableModel):
 def create_main_loop(tmpdir):
 
     def _create_main_loop(epochs=1, extra_hooks=(), dataset=None, model_class=None, skip_zeroth_epoch=True,
-                         **main_loop_kwargs):
+                         train_stream_name=EL_DEFAULT_TRAIN_STREAM, **main_loop_kwargs):
         """
         Create and return a model, dataset and mainloop.
 
@@ -291,6 +304,7 @@ def create_main_loop(tmpdir):
         :param dataset: dataset to be passed to the main loop, SimpleDataset() is created if None
         :param model_class: model class to be created and passed to the main loop, TrainableModel if None
         :param skip_zeroth_epoch: skip zeroth epoch flag passed to the main loop
+        :param train_stream_name: name of the training stream
         :return: a tuple of the created model, dataset and mainloop
         """
         hooks = list(extra_hooks) + [StopAfter(epochs=epochs)]
@@ -300,8 +314,8 @@ def create_main_loop(tmpdir):
             model_class = TrainableModel
         model = model_class(dataset=dataset, log_dir=tmpdir,  # pylint: disable=redefined-variable-type
                             io={'in': ['input', 'target'], 'out': ['output']})
-        mainloop = el.MainLoop(model=model, dataset=dataset, hooks=hooks,
-                               skip_zeroth_epoch=skip_zeroth_epoch, **main_loop_kwargs)
+        mainloop = el.MainLoop(model=model, dataset=dataset, hooks=hooks, skip_zeroth_epoch=skip_zeroth_epoch,
+                               train_stream_name=train_stream_name, **main_loop_kwargs)
         return model, dataset, mainloop
 
     return _create_main_loop
@@ -578,3 +592,31 @@ def test_stream_check(create_main_loop, caplog):
                                     'with `main_loop.fixed_size` = 47'),
             ('root', logging.INFO, 'Prediction done\n\n')
     ]
+
+
+def test_configurable_train_stream(create_main_loop, caplog):
+    caplog.set_level(logging.DEBUG)
+
+    caplog.clear()
+    _, _, mainloop = create_main_loop(dataset=EmptyStreamDataset(), train_stream_name='valid', on_empty_stream='warn')
+    mainloop.run_training()
+
+    assert caplog.record_tuples == [
+        ('root', logging.DEBUG, 'Training started'),
+        ('root', logging.INFO, 'Training epoch 1'),
+        ('root', logging.WARNING, 'Stream `valid` appears to be empty. Set `main_loop.on_empty_stream` to '
+                                  '`ignore` in order to suppress this warning.'),
+        ('root', logging.INFO, 'EpochStopperHook triggered'),
+        ('root', logging.INFO, 'Training terminated: Training terminated after epoch 1')
+    ]
+
+    caplog.clear()
+    _, _, mainloop = create_main_loop(dataset=GeneratedStreamDataset(), train_stream_name='generated')
+    mainloop.run_training()
+
+    assert caplog.record_tuples == [
+        ('root', logging.DEBUG, 'Training started'),
+        ('root', logging.INFO, 'Training epoch 1'),
+        ('root', logging.INFO, 'EpochStopperHook triggered'),
+        ('root', logging.INFO, 'Training terminated: Training terminated after epoch 1')
+    ]