Add reset_data_loader to SimpleTrainer

Summary: Add reset_data_loader( ) in SimpleTrainer for resetting the data loader.

Differential Revision: D37087154

fbshipit-source-id: 3d838214c50fd2e305ff55a8f2536fdd3494247f

detectron2/engine/train_loop.py

@@ -301,6 +301,11 @@ def _data_loader_iter(self):
             self._data_loader_iter_obj = iter(self.data_loader)
         return self._data_loader_iter_obj
 
+    def reset_data_loader(self, data_loader):
+        del self.data_loader
+        self.data_loader = data_loader
+        self._data_loader_iter_obj = None
+
     def _write_metrics(
         self,
         loss_dict: Mapping[str, torch.Tensor],

tests/test_engine.py

@@ -49,6 +49,15 @@ def test_simple_trainer(self, device="cpu"):
         )
         trainer.train(0, 10)
 
+    def test_simple_trainer_reset_dataloader(self, device="cpu"):
+        model = _SimpleModel().to(device=device)
+        trainer = SimpleTrainer(
+            model, self._data_loader(device), torch.optim.SGD(model.parameters(), 0.1)
+        )
+        trainer.train(0, 10)
+        trainer.reset_data_loader(self._data_loader(device))
+        trainer.train(0, 10)
+
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
     def test_simple_trainer_cuda(self):
         self.test_simple_trainer(device="cuda")