Repository: pytorch/ignite

Issue:
How to set a new dataloader during training?
## 🚀 Feature

Sometimes I want to replace my engine's dataloader with a new one during training. For example, it happens when I want to change a batch size or even change batch sampling strategy. Earlier, I used to change the field `engine.state.dataloader`, but it doesn't look correct, because there are some fields like `_dataloader_len` or `_dataloader_iter` that may need some adjustment as well.

Am I right about incorrectness of patching `engine.state.dataloader`? Is it possible to add something like `Engine.set_dataloader`?


Discussion:
@WeirdKeksButtonK thanks for the feedback! Yes, in the current version it would be a bit tricky to replace the dataloader just by replacing engine.state.dataloader. There are some other issue with the current implementation that we need to fix up. So, probably, we will also provide such helper method `set_dataloader`...

The most "proper" way (hack?) currently to replace the dataflow is to provide a dummy list to `run(data)`:
`data=list(range(total_num_iterations))` and set batch as required by your logic:
```python

dataloader = ...
next_dataloader = ...

def cycle(dataloader):
        while True:
            for i in dataloader:
                yield i

dataloader_iter = cycle(dataloader)

def udate_fn(engine, _):
    batch = next(dataloader_iter)
    ...

trainer = Engine(update_fn)

@trainer.on(Events.EPOCH_STARTED(once=50))
def switch_dataloader(_):
    dataloader_iter = cycle(next_dataloader)

data = list(range(len(dataloader) * num_epochs))
trainer.run(data, max_epochs=num_epochs)
```

HTH
@vfdev-5 thanks for the idea.
Indeed, it will be more convenient if updating the dataloader is supported natively.

And here is a version without dummy data:
```python
def _cycle(iterable):
    while True:
        for x in iterable:
            yield x


class DataLoaderBox:
    def __init__(self, dataloader):
        self.set_dataloader(dataloader)

    def __iter__(self):
        return self

    def __next__(self):
        return next(self._iter)

    def set_dataloader(self, dataloader):
        self._iter = _cycle(dataloader)


def update_fn(engine, batch):
    ...


dataloader = ...
next_dataloader = ...

trainer = Engine(update_fn)

@trainer.on(Events.EPOCH_STARTED(once=50))
def switch_dataloader(engine):
    engine.state.dataloader.set_dataloader(next_dataloader)

trainer.run(DataLoaderBox(dataloader))
```

It has its own drawbacks, but one may find it sufficient.

Resolving patch:
diff --git a/ignite/engine/engine.py b/ignite/engine/engine.py
--- a/ignite/engine/engine.py
+++ b/ignite/engine/engine.py
@@ -530,6 +530,42 @@ def load_state_dict(self, state_dict: Mapping) -> None:
     def _is_done(state: State) -> bool:
         return state.iteration == state.epoch_length * state.max_epochs
 
+    def set_data(self, data):
+        """Method to set data. After calling the method the next batch passed to `processing_function` is
+        from newly provided data. Please, note that epoch length is not modified.
+
+        Args:
+            data (Iterable): Collection of batches allowing repeated iteration (e.g., list or `DataLoader`).
+
+        Example usage:
+            User can switch data provider during the training:
+
+            .. code-block:: python
+
+                data1 = ...
+                data2 = ...
+
+                switch_iteration = 5000
+
+                def train_step(e, batch):
+                    # when iteration <= switch_iteration
+                    # batch is from data1
+                    # when iteration > switch_iteration
+                    # batch is from data2
+                    ...
+
+                trainer = Engine(train_step)
+
+                @trainer.on(Events.ITERATION_COMPLETED(once=switch_iteration))
+                def switch_dataloader():
+                    trainer.set_data(data2)
+
+                trainer.run(data1, max_epochs=100)
+
+        """
+        self.state.dataloader = data
+        self._dataloader_iter = iter(self.state.dataloader)
+
     def run(
         self,
         data: Iterable,
@@ -628,14 +664,6 @@ def _init_timers(state: State):
         state.times[Events.COMPLETED.name] = 0.0
 
     def _setup_engine(self) -> None:
-        try:
-            self._dataloader_len = None
-            if hasattr(self.state.dataloader, "__len__"):
-                self._dataloader_len = len(self.state.dataloader)
-        except TypeError:
-            # _InfiniteConstantSampler can raise a TypeError on DataLoader length of a IterableDataset
-            self._dataloader_len = None
-
         iteration = self.state.iteration
         self._dataloader_iter = iter(self.state.dataloader)
 
@@ -673,11 +701,11 @@ def _internal_run(self) -> State:
             self.logger.info("Engine run complete. Time taken %02d:%02d:%02d" % (hours, mins, secs))
 
         except BaseException as e:
-            self._dataloader_iter = self._dataloader_len = None
+            self._dataloader_iter = None
             self.logger.error("Engine run is terminating due to exception: %s.", str(e))
             self._handle_exception(e)
 
-        self._dataloader_iter = self._dataloader_len = None
+        self._dataloader_iter = None
         return self.state
 
     def _run_once_on_dataset(self) -> float:
@@ -697,14 +725,6 @@ def _run_once_on_dataset(self) -> float:
                     iter_counter += 1
                     should_exit = False
                 except StopIteration:
-
-                    if self._dataloader_len is None:
-                        if iter_counter > 0:
-                            self._dataloader_len = iter_counter
-                        else:
-                            # this can happen when data is finite iterator and epoch_length is equal to its size
-                            self._dataloader_len = self.state.iteration
-
                     # Define self.state.epoch_length if it is not yet set
                     if self.state.epoch_length is None:
                         # Define epoch length and stop the epoch
@@ -724,7 +744,7 @@ def _run_once_on_dataset(self) -> float:
                         break
 
                     self._fire_event(Events.DATALOADER_STOP_ITERATION)
-                    self._dataloader_iter = iter(self.state.dataloader)
+                    self.set_data(self.state.dataloader)
 
                     should_exit = True
 
@@ -741,7 +761,7 @@ def _run_once_on_dataset(self) -> float:
                 if self.should_terminate or self.should_terminate_single_epoch:
                     self._fire_event(Events.TERMINATE_SINGLE_EPOCH, iter_counter=iter_counter)
                     self.should_terminate_single_epoch = False
-                    self._dataloader_iter = iter(self.state.dataloader)
+                    self.set_data(self.state.dataloader)
                     break
 
                 if self.state.epoch_length is not None and iter_counter == self.state.epoch_length:


Test patch:
diff --git a/tests/ignite/engine/test_deterministic.py b/tests/ignite/engine/test_deterministic.py
--- a/tests/ignite/engine/test_deterministic.py
+++ b/tests/ignite/engine/test_deterministic.py
@@ -830,3 +830,26 @@ def foo(e, b):
     engine.run(data_loader, epoch_length=10, max_epochs=5)
 
     assert counter[0] == 50
+
+
+def test_run_finite_iterator_no_epoch_length():
+    # FR: https://github.com/pytorch/ignite/issues/871
+    unknown_size = 11
+
+    def finite_unk_size_data_iter():
+        for i in range(unknown_size):
+            yield i
+
+    bc = BatchChecker(data=list(range(unknown_size)))
+
+    engine = DeterministicEngine(lambda e, b: bc.check(b))
+
+    @engine.on(Events.DATALOADER_STOP_ITERATION)
+    def restart_iter():
+        engine.state.dataloader = finite_unk_size_data_iter()
+
+    data_iter = finite_unk_size_data_iter()
+    engine.run(data_iter, max_epochs=5)
+
+    assert engine.state.epoch == 5
+    assert engine.state.iteration == unknown_size * 5
diff --git a/tests/ignite/engine/test_engine.py b/tests/ignite/engine/test_engine.py
--- a/tests/ignite/engine/test_engine.py
+++ b/tests/ignite/engine/test_engine.py
@@ -332,7 +332,8 @@ def test__setup_engine():
     data = list(range(100))
     engine.state.dataloader = data
     engine._setup_engine()
-    assert engine._dataloader_len == len(data)
+    assert len(engine._init_iter) == 1 and engine._init_iter[0] == 10
+    # assert engine._dataloader_len == len(data)
 
 
 def test_run_asserts():
@@ -867,3 +868,31 @@ def val_step(evaluator, batch):
 
     assert evaluator.state.epoch == 1
     assert evaluator.state.iteration == size
+
+
+def test_set_data():
+    # tests FR https://github.com/pytorch/ignite/issues/833
+
+    num_iters1 = 10
+    num_iters2 = 20
+    batch_size = 4
+
+    torch.manual_seed(1)
+    data1 = torch.utils.data.DataLoader(torch.rand(num_iters1 * batch_size, 11), batch_size=batch_size)
+    data2 = torch.utils.data.DataLoader(torch.rand(num_iters2 * batch_size, 22), batch_size=batch_size)
+
+    switch_iteration = 35
+
+    def train_fn(e, batch):
+        if e.state.iteration <= switch_iteration:
+            assert batch.shape[1] == 11, "{}: {}".format(e.state.iteration, batch.shape)
+        else:
+            assert batch.shape[1] == 22, "{}: {}".format(e.state.iteration, batch.shape)
+
+    trainer = Engine(train_fn)
+
+    @trainer.on(Events.ITERATION_COMPLETED(once=switch_iteration))
+    def switch_dataloader():
+        trainer.set_data(data2)
+
+    trainer.run(data1, max_epochs=10)