[Enhancement] Support overall dastaloader settings

open-mmlab · Apr 13, 2022 · ee7a689 · ee7a689
1 parent 31568b1
commit ee7a689
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 4 deletions.
diff --git a/mmflow/apis/train.py b/mmflow/apis/train.py
@@ -96,14 +96,25 @@ def train_model(model: Module,
 
     # prepare data loaders
     dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
+
+    # The overall dataloader settings
+    loader_cfg = {
+        k: v
+        for k, v in cfg.data.items() if k not in [
+            'train', 'val', 'test', 'train_dataloader', 'val_dataloader',
+            'test_dataloader'
+        ]
+    }
+    # The specific training dataloader settings
+    train_loader_cfg = {**loader_cfg, **cfg.data.get('train_dataloader', {})}
     data_loaders = [
         build_dataloader(
             ds,
             # cfg.gpus will be ignored if distributed
             num_gpus=len(cfg.gpu_ids),
             dist=distributed,
             seed=cfg.seed,
-            **cfg.data.train_dataloader) for ds in dataset
+            **train_loader_cfg) for ds in dataset
     ]
 
     # put model on gpus
@@ -160,6 +171,9 @@ def train_model(model: Module,
     else:
         optimizer_config = cfg.optimizer_config
 
+    # The specific validation dataloader settings
+    val_loader_cfg = {**loader_cfg, **cfg.data.get('val_dataloader', {})}
+
     # register hooks
     runner.register_training_hooks(cfg.lr_config, optimizer_config,
                                    cfg.checkpoint_config, cfg.log_config,
@@ -174,15 +188,15 @@ def train_model(model: Module,
             ]
             val_dataloader = [
                 build_dataloader(
-                    _val_dataset, **cfg.data.val_dataloader, dist=distributed)
+                    _val_dataset, **val_loader_cfg, dist=distributed)
                 for _val_dataset in val_dataset
             ]
             val_dataset_name = [ds.dataset_name for ds in val_dataset]
         else:
 
             val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
             val_dataloader = build_dataloader(
-                val_dataset, **cfg.data.val_dataloader, dist=distributed)
+                val_dataset, **val_loader_cfg, dist=distributed)
             val_dataset_name = val_dataset.dataset_name
 
         eval_cfg = cfg.get('evaluation', {})

diff --git a/tools/test.py b/tools/test.py
@@ -120,6 +120,17 @@ def main():
     # set multi-process settings
     setup_multi_processes(cfg)
 
+    # The overall dataloader settings
+    loader_cfg = {
+        k: v
+        for k, v in cfg.data.items() if k not in [
+            'train', 'val', 'test', 'train_dataloader', 'val_dataloader',
+            'test_dataloader'
+        ]
+    }
+    # The specific training dataloader settings
+    test_loader_cfg = {**loader_cfg, **cfg.data.get('test_dataloader', {})}
+
     # build the dataloader
     separate_eval = cfg.data.test.get('separate_eval', False)
     if separate_eval:
@@ -133,7 +144,7 @@ def main():
     data_loader = [
         build_dataloader(
             _dataset,
-            **cfg.data.test_dataloader,
+            **test_loader_cfg,
             dist=distributed,
         ) for _dataset in dataset
     ]