use public trainer API to re-create data loader for QAT

Summary:
Pull Request resolved: facebookresearch#4458

X-link: facebookresearch/d2go#353

- The QAT was using old code prior to D36786902 (facebookresearch@e566079), update to use public API
- Make `trainer:reset_data_loader` to take lazy lambda expression in order to delay the creation of dataloader. It's possible that we don't have enough RAM to hold two data loader at the same time, so we need to delete the first one, then create the second one.

Differential Revision: D38330148

fbshipit-source-id: aae28a48eabf211fe00cafe5d9ea8aeaf56e4e0c

detectron2/engine/train_loop.py

@@ -301,8 +301,13 @@ 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):
+    def reset_data_loader(self, data_loader_builder):
+        """
+        Delete and replace the current data loader with a new one, which will be created
+        by calling `data_loader_builder` (without argument).
+        """
         del self.data_loader
+        data_loader = data_loader_builder()
         self.data_loader = data_loader
         self._data_loader_iter_obj = None
 

tests/test_engine.py

@@ -55,7 +55,7 @@ def test_simple_trainer_reset_dataloader(self, device="cpu"):
             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.reset_data_loader(lambda: self._data_loader(device))
         trainer.train(0, 10)
 
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")