From 88d4a8ddb1a22efd86d69025c9dd3764a0ff8a6c Mon Sep 17 00:00:00 2001 From: wutongshenqiu <44188071+wutongshenqiu@users.noreply.github.com> Date: Wed, 13 Apr 2022 14:34:37 +0800 Subject: [PATCH] add mmcls and mmseg --- mmrazor/apis/mmcls/train.py | 56 ++++++++++++++++++------------------- mmrazor/apis/mmseg/train.py | 44 ++++++++++++++++++----------- tools/mmcls/test_mmcls.py | 33 ++++++++++++++++------ tools/mmseg/test_mmseg.py | 24 +++++++++++++--- 4 files changed, 100 insertions(+), 57 deletions(-) diff --git a/mmrazor/apis/mmcls/train.py b/mmrazor/apis/mmcls/train.py index 91aafc4bf..89e2a5c50 100644 --- a/mmrazor/apis/mmcls/train.py +++ b/mmrazor/apis/mmcls/train.py @@ -65,7 +65,24 @@ def train_mmcls_model(model, train_dataset = dataset[0] dataset[0] = split_dataset(train_dataset) - sampler_cfg = cfg.data.get('sampler', None) + loader_cfg = dict( + # cfg.gpus will be ignored if distributed + num_gpus=len(cfg.gpu_ids), + dist=distributed, + round_up=True, + seed=cfg.get('seed'), + sampler_cfg=cfg.get('sampler', None), + ) + # The overall dataloader settings + loader_cfg.update({ + k: v + for k, v in cfg.data.items() if k not in [ + 'train', 'val', 'test', 'train_dataloader', 'val_dataloader', + 'test_dataloader' + ] + }) + # The specific dataloader settings + train_loader_cfg = {**loader_cfg, **cfg.data.get('train_dataloader', {})} # Difference from mmclassification. # Build multi dataloaders according the splited datasets. @@ -73,28 +90,11 @@ def train_mmcls_model(model, for dset in dataset: if isinstance(dset, list): data_loader = [ - build_dataloader( - item_ds, - cfg.data.samples_per_gpu, - cfg.data.workers_per_gpu, - # cfg.gpus will be ignored if distributed - num_gpus=len(cfg.gpu_ids), - dist=distributed, - round_up=True, - seed=cfg.seed, - sampler_cfg=sampler_cfg) for item_ds in dset + build_dataloader(item_ds, **train_loader_cfg) + for item_ds in dset ] else: - data_loader = build_dataloader( - dset, - cfg.data.samples_per_gpu, - cfg.data.workers_per_gpu, - # cfg.gpus will be ignored if distributed - num_gpus=len(cfg.gpu_ids), - dist=distributed, - round_up=True, - seed=cfg.seed, - sampler_cfg=sampler_cfg) + data_loader = build_dataloader(dset, **train_loader_cfg) data_loaders.append(data_loader) @@ -188,13 +188,13 @@ def train_mmcls_model(model, # register eval hooks if validate: val_dataset = build_dataset(cfg.data.val, dict(test_mode=True)) - val_dataloader = build_dataloader( - val_dataset, - samples_per_gpu=cfg.data.samples_per_gpu, - workers_per_gpu=cfg.data.workers_per_gpu, - dist=distributed, - shuffle=False, - round_up=True) + val_loader_cfg = { + **loader_cfg, + 'shuffle': False, # Not shuffle by default + 'sampler_cfg': None, # Not use sampler by default + **cfg.data.get('val_dataloader', {}), + } + val_dataloader = build_dataloader(val_dataset, **val_loader_cfg) eval_cfg = cfg.get('evaluation', {}) eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner' diff --git a/mmrazor/apis/mmseg/train.py b/mmrazor/apis/mmseg/train.py index 16de892f9..ed09df78e 100644 --- a/mmrazor/apis/mmseg/train.py +++ b/mmrazor/apis/mmseg/train.py @@ -57,17 +57,25 @@ def train_mmseg_model(model, # prepare data loaders dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset] - data_loaders = [ - build_dataloader( - ds, - cfg.data.samples_per_gpu, - cfg.data.workers_per_gpu, - # cfg.gpus will be ignored if distributed - len(cfg.gpu_ids), - dist=distributed, - seed=cfg.seed, - drop_last=True) for ds in dataset - ] + # The default loader config + loader_cfg = dict( + # cfg.gpus will be ignored if distributed + num_gpus=len(cfg.gpu_ids), + dist=distributed, + seed=cfg.seed, + drop_last=True) + # The overall dataloader settings + loader_cfg.update({ + k: v + for k, v in cfg.data.items() if k not in [ + 'train', 'val', 'test', 'train_dataloader', 'val_dataloader', + 'test_dataloader' + ] + }) + + # The specific dataloader settings + train_loader_cfg = {**loader_cfg, **cfg.data.get('train_dataloader', {})} + data_loaders = [build_dataloader(ds, **train_loader_cfg) for ds in dataset] # put model on gpus if distributed: @@ -130,12 +138,14 @@ def train_mmseg_model(model, # register eval hooks if validate: val_dataset = build_dataset(cfg.data.val, dict(test_mode=True)) - val_dataloader = build_dataloader( - val_dataset, - samples_per_gpu=1, - workers_per_gpu=cfg.data.workers_per_gpu, - dist=distributed, - shuffle=False) + # The specific dataloader settings + val_loader_cfg = { + **loader_cfg, + 'samples_per_gpu': 1, + 'shuffle': False, # Not shuffle by default + **cfg.data.get('val_dataloader', {}), + } + val_dataloader = build_dataloader(val_dataset, **val_loader_cfg) eval_cfg = cfg.get('evaluation', {}) eval_cfg['by_epoch'] = cfg.runner['type'] != 'IterBasedRunner' eval_hook = DistEvalHook if distributed else EvalHook diff --git a/tools/mmcls/test_mmcls.py b/tools/mmcls/test_mmcls.py index 8ee6351f4..26b55077f 100644 --- a/tools/mmcls/test_mmcls.py +++ b/tools/mmcls/test_mmcls.py @@ -120,15 +120,32 @@ def main(): init_dist(args.launcher, **cfg.dist_params) # build the dataloader - dataset = build_dataset(cfg.data.test) - # the extra round_up data will be removed during gpu/cpu collect - data_loader = build_dataloader( - dataset, - samples_per_gpu=cfg.data.samples_per_gpu, - workers_per_gpu=cfg.data.workers_per_gpu, + dataset = build_dataset(cfg.data.test, default_args=dict(test_mode=True)) + + # build the dataloader + # The default loader config + loader_cfg = dict( + # cfg.gpus will be ignored if distributed + num_gpus=len(cfg.gpu_ids), dist=distributed, - shuffle=False, - round_up=True) + round_up=True, + ) + # The overall dataloader settings + loader_cfg.update({ + k: v + for k, v in cfg.data.items() if k not in [ + 'train', 'val', 'test', 'train_dataloader', 'val_dataloader', + 'test_dataloader' + ] + }) + test_loader_cfg = { + **loader_cfg, + 'shuffle': False, # Not shuffle by default + 'sampler_cfg': None, # Not use sampler by default + **cfg.data.get('test_dataloader', {}), + } + # the extra round_up data will be removed during gpu/cpu collect + data_loader = build_dataloader(dataset, **test_loader_cfg) # build the algorithm and load checkpoint algorithm = build_algorithm(cfg.algorithm) diff --git a/tools/mmseg/test_mmseg.py b/tools/mmseg/test_mmseg.py index 3528ff8d1..50a06eb4e 100644 --- a/tools/mmseg/test_mmseg.py +++ b/tools/mmseg/test_mmseg.py @@ -154,12 +154,28 @@ def main(): # build the dataloader # TODO: support multiple images per gpu (only minor changes are needed) dataset = build_dataset(cfg.data.test) - data_loader = build_dataloader( - dataset, - samples_per_gpu=1, - workers_per_gpu=cfg.data.workers_per_gpu, + # The default loader config + loader_cfg = dict( + # cfg.gpus will be ignored if distributed + num_gpus=len(cfg.gpu_ids), dist=distributed, shuffle=False) + # The overall dataloader settings + loader_cfg.update({ + k: v + for k, v in cfg.data.items() if k not in [ + 'train', 'val', 'test', 'train_dataloader', 'val_dataloader', + 'test_dataloader' + ] + }) + test_loader_cfg = { + **loader_cfg, + 'samples_per_gpu': 1, + 'shuffle': False, # Not shuffle by default + **cfg.data.get('test_dataloader', {}) + } + # build the dataloader + data_loader = build_dataloader(dataset, **test_loader_cfg) # build the algorithm and load checkpoint # Difference from mmsegmentation