Make InOrder iterator stateful (#993)

Summary: Pull Request resolved: #993

Reviewed By: galrotem

Differential Revision: D73863744

fbshipit-source-id: 42eaec46d64945fa0a481fda98a4c15a40a8239b

Author: diego-urgell

tests/utils/data/test_multi_dataloader.py

@@ -14,12 +14,16 @@
 
 import torch
 from torch.utils.data import DataLoader, Dataset
-from torchtnt.framework._test_utils import generate_random_dataloader
+from torchtnt.framework._test_utils import (
+    generate_random_dataloader,
+    generate_tensor_dataloader,
+)
 
 from torchtnt.utils.data.iterators import (
     AllDatasetBatches,
     DataIterationStrategy,
     InOrder,
+    InOrderIterator,
     MultiIterator,
     RandomizedBatchSampler,
     RoundRobin,
@@ -580,3 +584,68 @@ def _get_dataloaders_dict(
         self, first_dataloader: DataLoader, second_dataloader: DataLoader
     ) -> Dict[str, Union[DataLoader, Iterable[object]]]:
         return {"1": first_dataloader, "2": second_dataloader}
+
+    def test_state_dict_with_inorder_iterator(self) -> None:
+        dataloader_1 = generate_tensor_dataloader(torch.tensor([1, 2]), batch_size=1)
+        dataloader_2 = generate_tensor_dataloader(torch.tensor([3, 4]), batch_size=1)
+        dataloader_3 = generate_tensor_dataloader(torch.tensor([5, 6]), batch_size=1)
+
+        dataloaders_dict: Dict[str, Union[DataLoader, Iterable[object]]] = {
+            "dataloader_1": dataloader_1,
+            "dataloader_2": dataloader_2,
+            "dataloader_3": dataloader_3,
+        }
+
+        multi_dataloader = MultiDataLoader(
+            dataloaders_dict,
+            InOrder(iteration_order=["dataloader_1", "dataloader_2", "dataloader_3"]),
+        )
+
+        mdl_iter = iter(multi_dataloader)
+
+        # Exhaust first dataset
+        self.assertEqual(next(mdl_iter)["dataloader_1"], [torch.tensor([1])])
+        self.assertEqual(next(mdl_iter)["dataloader_1"], [torch.tensor([2])])
+
+        # We expect same iterator since we haven't raise StopIteration
+        mdl_sd = multi_dataloader.state_dict()
+        self.assertEqual(
+            mdl_sd["iterator_state"],
+            {
+                "iterators_finished": 0,
+                "cur_iterator": "dataloader_1",
+            },
+        )
+
+        # Start second dataset
+        self.assertEqual(next(mdl_iter)["dataloader_2"], [torch.tensor([3])])
+
+        mdl_sd = multi_dataloader.state_dict()
+        self.assertEqual(
+            mdl_sd["iterator_state"],
+            {
+                "iterators_finished": 1,
+                "cur_iterator": "dataloader_2",
+            },
+        )
+
+        # Create new dataloader and verify restore
+        multi_dataloader_2 = MultiDataLoader(
+            dataloaders_dict,
+            InOrder(iteration_order=["dataloader_1", "dataloader_2", "dataloader_3"]),
+        )
+        multi_dataloader_2.load_state_dict(mdl_sd)
+        in_order_iter = cast(InOrderIterator, iter(multi_dataloader_2))
+        self.assertEqual(in_order_iter.cur_iterator, "dataloader_1")
+        self.assertEqual(in_order_iter.iterators_finished, 1)
+
+        # Calling next should update the currrent iterator
+        # individual dl is not stateful
+        self.assertEqual(next(in_order_iter)["dataloader_2"], [torch.tensor([3])])
+        self.assertEqual(in_order_iter.cur_iterator, "dataloader_2")
+
+        # verify that after calling iter(), everything is reset
+        in_order_iter_2 = cast(InOrderIterator, iter(multi_dataloader_2))
+        self.assertEqual(in_order_iter_2.cur_iterator, "dataloader_1")
+        self.assertEqual(in_order_iter_2.iterators_finished, 0)
+        self.assertEqual(next(in_order_iter_2)["dataloader_1"], [torch.tensor([1])])

torchtnt/framework/_test_utils.py

@@ -280,6 +280,17 @@ def generate_dummy_stateful_dataloader(
     )
 
 
+def generate_tensor_dataloader(
+    samples: torch.Tensor, batch_size: int
+) -> DummyStatefulDataLoader:
+    return DummyStatefulDataLoader(
+        DataLoader(
+            dataset=TensorDataset(samples),
+            batch_size=batch_size,
+        )
+    )
+
+
 class DummyMeanMetric:
     def __init__(self) -> None:
         super().__init__()

torchtnt/utils/data/iterators.py

@@ -511,6 +511,7 @@ def __init__(
         self.cur_iter: Union[Iterator[DataLoader], Iterator[object]] = iter(
             self.individual_dataloaders[self.iteration_order[0]]
         )
+        self.cur_iterator_idx: int = 0
         self.cur_iterator: str = self.iteration_order[0]
         self.num_iterators: int = len(self.iteration_order)
         self.iterators_finished: int = 0
@@ -519,6 +520,17 @@ def __next__(self) -> Dict[str, Any]:
         if self.iterators_finished == self.num_iterators:
             raise StopIteration
 
+        # If the current iterator doesn't match the expected number of finished iterators,
+        # it means we restored from checkpoint and we need to initialize expected iterator
+        # This is to avoid calling iter() in the load_state_dict() function.
+        if self.iterators_finished != self.cur_iterator_idx:
+            logger.info(
+                f"Initializing iterator {self.cur_iterator} after resuming from checkpoint"
+            )
+            self.cur_iterator_idx = self.iterators_finished
+            self.cur_iterator = self.iteration_order[self.iterators_finished]
+            self.cur_iter = iter(self.individual_dataloaders[self.cur_iterator])
+
         try:
             return {self.cur_iterator: next(self.cur_iter)}
         except StopIteration:
@@ -528,12 +540,36 @@ def __next__(self) -> Dict[str, Any]:
             if self.iterators_finished == self.num_iterators:
                 raise StopIteration
 
+            self.cur_iterator_idx += 1
             self.cur_iterator = self.iteration_order[self.iterators_finished]
-
             self.cur_iter = iter(self.individual_dataloaders[self.cur_iterator])
 
             return self.__next__()
 
+    def state_dict(self) -> Dict[str, Any]:
+        return {
+            "iterators_finished": self.iterators_finished,
+            "cur_iterator": self.cur_iterator,
+        }
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        iterators_finished: int = state_dict["iterators_finished"]
+        cur_iterator: str = state_dict["cur_iterator"]
+        logger.info(
+            f"Loading InOrderIterator state. Trying to set iterators_finished to {iterators_finished} to restore {cur_iterator}"
+        )
+
+        if cur_iterator not in self.iteration_order or iterators_finished > len(
+            self.iteration_order
+        ):
+            logger.warning(
+                f"Will not restore InOrderIterator state, since expected dataloader was not found in available iterators: {cur_iterator}"
+            )
+            return
+
+        self.iterators_finished = iterators_finished
+        # We do not initialize actual iterator here to avoid checkpoint restore taking longer
+
 
 class DataIterationStrategyRegistry:
     """A generic iterator registry.