We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
a dataset bug causing topdown training very slow, wasting 3 min every epoch
i found a dataset bug, i test it on several server(including 8 a100 with 96 core cpu), it all happened. for every epoch, this bug cause about 3min time wasting. i jsut can locat the bug, but i don't known why it happen. it seems only happen when distribution launching.

bug loaction: when you lauch a topdown method, eg, topdown_heatmap/coco/res50_coco_256x192.py, go to /mmcv/runner/epoch_based_runner.py, about line 48. there is such func

        self.call_hook('before_train_epoch')
        time.sleep(2)  # Prevent possible deadlock during epoch transition
        for i, data_batch in enumerate(self.data_loader):
            self._inner_iter = i

at the every epoch begining, the ( for i, data_batch in enumerate(self.data_loader): ) takes about 3min, it make the training very slow.

you can modify the ori code to the code below to reproduce this issue, this only happen at very epoch begining.

        self.call_hook('before_train_epoch')
        time.sleep(2)  # Prevent possible deadlock during epoch transition
        print('before_train_epoch, time:{}'.format(time.time()-start_time))
        start_time = time.time()
        for i, data_batch in enumerate(self.data_loader):
            self._inner_iter = i
            print('before_train_iter_load_data, time:{}'.format(time.time()-start_time))

here is my sys information
Python: 3.7.10 | packaged by conda-forge | (default, Feb 19 2021, 16:07:37) [GCC 9.3.0]                                                                                       
CUDA available: True                                                                                                                                                          GPU 0,1,2,3,4,5,6,7: A100-SXM4-40GB                                                                                                                                           
CUDA_HOME: /usr/local/cuda-11.1                                                                                                                                               
NVCC: Build cuda_11.1.TC455_06.29190527_0                                                                                                                                     GCC: gcc (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0                                                                                                                                  
PyTorch: 1.8.1+cu111                                                                                                                                                          
PyTorch compiling details: PyTorch built with:                                                                                                                                
  - GCC 7.3                    
  - C++ Version: 201402            
  - Intel(R) Math Kernel Library Version 2020.0.0 Product Build 20191122 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v1.7.0 (Git Hash 7aed236906b1f7a05c0917e5257a1af05e9ff683)
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - NNPACK is enabled                                           
  - CPU capability usage: AVX2                  
  - CUDA Runtime 11.1
  - NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_70,code=sm_70;-
gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86
  - CuDNN 8.0.5                                                             
  - Magma 2.5.2                            
  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=11.1, CUDNN_VERSION=8.0.5, CXX_COMPILER=/opt/rh/devtoolset-7/root/usr/bin/c++, CXX_FLAGS= -Wno-deprecated
-fvisibility-inlines-hidden -DUSE_PTHREADPOOL -fopenmp -DNDEBUG -DUSE_KINETO -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -O2 -fPIC -Wno-narrowing -Wall -We
xtra -Werror=return-type -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-sign-compare -Wno-unused-parameter -Wno-unused-variable
-Wno-unused-function -Wno-unused-result -Wno-unused-local-typedefs -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-stringop-overflow -Wno-ps
abi -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -f
no-math-errno -fno-trapping-math -Werror=format -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_VERSION=1.8.1, USE_CUDA=
ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON,
                                   
TorchVision: 0.9.1+cu111  
OpenCV: 4.5.3             
MMCV: 1.3.8                
MMCV Compiler: GCC 7.5
MMCV CUDA Compiler: 11.1
MMPose: 0.15.0+51b4b45

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mmpose/apis/train.py`
Content:
```
1 # Copyright (c) OpenMMLab. All rights reserved.
2 import warnings
3 
4 import torch
5 from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
6 from mmcv.runner import DistSamplerSeedHook, EpochBasedRunner, OptimizerHook
7 
8 from mmpose.core import DistEvalHook, EvalHook, build_optimizers
9 from mmpose.core.distributed_wrapper import DistributedDataParallelWrapper
10 from mmpose.datasets import build_dataloader, build_dataset
11 from mmpose.utils import get_root_logger
12 
13 try:
14     from mmcv.runner import Fp16OptimizerHook
15 except ImportError:
16     warnings.warn(
17         'Fp16OptimizerHook from mmpose will be deprecated from '
18         'v0.15.0. Please install mmcv>=1.1.4', DeprecationWarning)
19     from mmpose.core import Fp16OptimizerHook
20 
21 
22 def train_model(model,
23                 dataset,
24                 cfg,
25                 distributed=False,
26                 validate=False,
27                 timestamp=None,
28                 meta=None):
29     """Train model entry function.
30 
31     Args:
32         model (nn.Module): The model to be trained.
33         dataset (Dataset): Train dataset.
34         cfg (dict): The config dict for training.
35         distributed (bool): Whether to use distributed training.
36             Default: False.
37         validate (bool): Whether to do evaluation. Default: False.
38         timestamp (str | None): Local time for runner. Default: None.
39         meta (dict | None): Meta dict to record some important information.
40             Default: None
41     """
42     logger = get_root_logger(cfg.log_level)
43 
44     # prepare data loaders
45     dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
46     dataloader_setting = dict(
47         samples_per_gpu=cfg.data.get('samples_per_gpu', {}),
48         workers_per_gpu=cfg.data.get('workers_per_gpu', {}),
49         # cfg.gpus will be ignored if distributed
50         num_gpus=len(cfg.gpu_ids),
51         dist=distributed,
52         seed=cfg.seed)
53     dataloader_setting = dict(dataloader_setting,
54                               **cfg.data.get('train_dataloader', {}))
55 
56     data_loaders = [
57         build_dataloader(ds, **dataloader_setting) for ds in dataset
58     ]
59 
60     # determine wether use adversarial training precess or not
61     use_adverserial_train = cfg.get('use_adversarial_train', False)
62 
63     # put model on gpus
64     if distributed:
65         find_unused_parameters = cfg.get('find_unused_parameters', True)
66         # Sets the `find_unused_parameters` parameter in
67         # torch.nn.parallel.DistributedDataParallel
68 
69         if use_adverserial_train:
70             # Use DistributedDataParallelWrapper for adversarial training
71             model = DistributedDataParallelWrapper(
72                 model,
73                 device_ids=[torch.cuda.current_device()],
74                 broadcast_buffers=False,
75                 find_unused_parameters=find_unused_parameters)
76         else:
77             model = MMDistributedDataParallel(
78                 model.cuda(),
79                 device_ids=[torch.cuda.current_device()],
80                 broadcast_buffers=False,
81                 find_unused_parameters=find_unused_parameters)
82     else:
83         model = MMDataParallel(
84             model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)
85 
86     # build runner
87     optimizer = build_optimizers(model, cfg.optimizer)
88 
89     runner = EpochBasedRunner(
90         model,
91         optimizer=optimizer,
92         work_dir=cfg.work_dir,
93         logger=logger,
94         meta=meta)
95     # an ugly workaround to make .log and .log.json filenames the same
96     runner.timestamp = timestamp
97 
98     if use_adverserial_train:
99         # The optimizer step process is included in the train_step function
100         # of the model, so the runner should NOT include optimizer hook.
101         optimizer_config = None
102     else:
103         # fp16 setting
104         fp16_cfg = cfg.get('fp16', None)
105         if fp16_cfg is not None:
106             optimizer_config = Fp16OptimizerHook(
107                 **cfg.optimizer_config, **fp16_cfg, distributed=distributed)
108         elif distributed and 'type' not in cfg.optimizer_config:
109             optimizer_config = OptimizerHook(**cfg.optimizer_config)
110         else:
111             optimizer_config = cfg.optimizer_config
112 
113     # register hooks
114     runner.register_training_hooks(cfg.lr_config, optimizer_config,
115                                    cfg.checkpoint_config, cfg.log_config,
116                                    cfg.get('momentum_config', None))
117     if distributed:
118         runner.register_hook(DistSamplerSeedHook())
119 
120     # register eval hooks
121     if validate:
122         eval_cfg = cfg.get('evaluation', {})
123         val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
124         dataloader_setting = dict(
125             samples_per_gpu=1,
126             workers_per_gpu=cfg.data.get('workers_per_gpu', {}),
127             # cfg.gpus will be ignored if distributed
128             num_gpus=len(cfg.gpu_ids),
129             dist=distributed,
130             drop_last=False,
131             shuffle=False)
132         dataloader_setting = dict(dataloader_setting,
133                                   **cfg.data.get('val_dataloader', {}))
134         val_dataloader = build_dataloader(val_dataset, **dataloader_setting)
135         eval_hook = DistEvalHook if distributed else EvalHook
136         runner.register_hook(eval_hook(val_dataloader, **eval_cfg))
137 
138     if cfg.resume_from:
139         runner.resume(cfg.resume_from)
140     elif cfg.load_from:
141         runner.load_checkpoint(cfg.load_from)
142     runner.run(data_loaders, cfg.workflow, cfg.total_epochs)
143
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mmpose/apis/train.py b/mmpose/apis/train.py
--- a/mmpose/apis/train.py
+++ b/mmpose/apis/train.py
@@ -43,19 +43,33 @@
 
     # prepare data loaders
     dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
-    dataloader_setting = dict(
-        samples_per_gpu=cfg.data.get('samples_per_gpu', {}),
-        workers_per_gpu=cfg.data.get('workers_per_gpu', {}),
-        # cfg.gpus will be ignored if distributed
-        num_gpus=len(cfg.gpu_ids),
-        dist=distributed,
-        seed=cfg.seed)
-    dataloader_setting = dict(dataloader_setting,
-                              **cfg.data.get('train_dataloader', {}))
-
-    data_loaders = [
-        build_dataloader(ds, **dataloader_setting) for ds in dataset
-    ]
+    # step 1: give default values and override (if exist) from cfg.data
+    loader_cfg = {
+        **dict(
+            seed=cfg.get('seed'),
+            drop_last=False,
+            dist=distributed,
+            num_gpus=len(cfg.gpu_ids)),
+        **({} if torch.__version__ != 'parrots' else dict(
+               prefetch_num=2,
+               pin_memory=False,
+           )),
+        **dict((k, cfg.data[k]) for k in [
+                   'samples_per_gpu',
+                   'workers_per_gpu',
+                   'shuffle',
+                   'seed',
+                   'drop_last',
+                   'prefetch_num',
+                   'pin_memory',
+                   'persistent_workers',
+               ] if k in cfg.data)
+    }
+
+    # step 2: cfg.data.train_dataloader has highest priority
+    train_loader_cfg = dict(loader_cfg, **cfg.data.get('train_dataloader', {}))
+
+    data_loaders = [build_dataloader(ds, **train_loader_cfg) for ds in dataset]
 
     # determine wether use adversarial training precess or not
     use_adverserial_train = cfg.get('use_adversarial_train', False)

{"golden_diff": "diff --git a/mmpose/apis/train.py b/mmpose/apis/train.py\n--- a/mmpose/apis/train.py\n+++ b/mmpose/apis/train.py\n@@ -43,19 +43,33 @@\n \n     # prepare data loaders\n     dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]\n-    dataloader_setting = dict(\n-        samples_per_gpu=cfg.data.get('samples_per_gpu', {}),\n-        workers_per_gpu=cfg.data.get('workers_per_gpu', {}),\n-        # cfg.gpus will be ignored if distributed\n-        num_gpus=len(cfg.gpu_ids),\n-        dist=distributed,\n-        seed=cfg.seed)\n-    dataloader_setting = dict(dataloader_setting,\n-                              **cfg.data.get('train_dataloader', {}))\n-\n-    data_loaders = [\n-        build_dataloader(ds, **dataloader_setting) for ds in dataset\n-    ]\n+    # step 1: give default values and override (if exist) from cfg.data\n+    loader_cfg = {\n+        **dict(\n+            seed=cfg.get('seed'),\n+            drop_last=False,\n+            dist=distributed,\n+            num_gpus=len(cfg.gpu_ids)),\n+        **({} if torch.__version__ != 'parrots' else dict(\n+               prefetch_num=2,\n+               pin_memory=False,\n+           )),\n+        **dict((k, cfg.data[k]) for k in [\n+                   'samples_per_gpu',\n+                   'workers_per_gpu',\n+                   'shuffle',\n+                   'seed',\n+                   'drop_last',\n+                   'prefetch_num',\n+                   'pin_memory',\n+                   'persistent_workers',\n+               ] if k in cfg.data)\n+    }\n+\n+    # step 2: cfg.data.train_dataloader has highest priority\n+    train_loader_cfg = dict(loader_cfg, **cfg.data.get('train_dataloader', {}))\n+\n+    data_loaders = [build_dataloader(ds, **train_loader_cfg) for ds in dataset]\n \n     # determine wether use adversarial training precess or not\n     use_adverserial_train = cfg.get('use_adversarial_train', False)\n", "issue": "a dataset bug causing topdown training very slow, wasting 3 min every epoch\ni found a dataset bug, i test it on several server(including 8 a100 with 96 core cpu), it all happened. for every epoch, this bug cause about 3min time wasting. i jsut can locat the bug, but i don't known why it happen. it seems only happen when distribution launching.\r\n\r\nbug loaction: when you lauch a topdown method, eg, topdown_heatmap/coco/res50_coco_256x192.py, go to /mmcv/runner/epoch_based_runner.py, about line 48. there is such func\r\n\r\n        self.call_hook('before_train_epoch')\r\n        time.sleep(2)  # Prevent possible deadlock during epoch transition\r\n        for i, data_batch in enumerate(self.data_loader):\r\n            self._inner_iter = i\r\n\r\nat the every epoch begining, the ( for i, data_batch in enumerate(self.data_loader): ) takes about 3min, it make the training very slow.\r\n\r\nyou can modify the ori code to the code below to reproduce this issue, this only happen at very epoch begining.\r\n\r\n        self.call_hook('before_train_epoch')\r\n        time.sleep(2)  # Prevent possible deadlock during epoch transition\r\n        print('before_train_epoch, time:{}'.format(time.time()-start_time))\r\n        start_time = time.time()\r\n        for i, data_batch in enumerate(self.data_loader):\r\n            self._inner_iter = i\r\n            print('before_train_iter_load_data, time:{}'.format(time.time()-start_time))\r\n\r\nhere is my sys information\r\nPython: 3.7.10 | packaged by conda-forge | (default, Feb 19 2021, 16:07:37) [GCC 9.3.0]                                                                                       \r\nCUDA available: True                                                                                                                                                          GPU 0,1,2,3,4,5,6,7: A100-SXM4-40GB                                                                                                                                           \r\nCUDA_HOME: /usr/local/cuda-11.1                                                                                                                                               \r\nNVCC: Build cuda_11.1.TC455_06.29190527_0                                                                                                                                     GCC: gcc (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0                                                                                                                                  \r\nPyTorch: 1.8.1+cu111                                                                                                                                                          \r\nPyTorch compiling details: PyTorch built with:                                                                                                                                \r\n  - GCC 7.3                    \r\n  - C++ Version: 201402            \r\n  - Intel(R) Math Kernel Library Version 2020.0.0 Product Build 20191122 for Intel(R) 64 architecture applications\r\n  - Intel(R) MKL-DNN v1.7.0 (Git Hash 7aed236906b1f7a05c0917e5257a1af05e9ff683)\r\n  - OpenMP 201511 (a.k.a. OpenMP 4.5)\r\n  - NNPACK is enabled                                           \r\n  - CPU capability usage: AVX2                  \r\n  - CUDA Runtime 11.1\r\n  - NVCC architecture flags: -gencode;arch=compute_37,code=sm_37;-gencode;arch=compute_50,code=sm_50;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_70,code=sm_70;-\r\ngencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86\r\n  - CuDNN 8.0.5                                                             \r\n  - Magma 2.5.2                            \r\n  - Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=11.1, CUDNN_VERSION=8.0.5, CXX_COMPILER=/opt/rh/devtoolset-7/root/usr/bin/c++, CXX_FLAGS= -Wno-deprecated\r\n-fvisibility-inlines-hidden -DUSE_PTHREADPOOL -fopenmp -DNDEBUG -DUSE_KINETO -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -O2 -fPIC -Wno-narrowing -Wall -We\r\nxtra -Werror=return-type -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-sign-compare -Wno-unused-parameter -Wno-unused-variable\r\n-Wno-unused-function -Wno-unused-result -Wno-unused-local-typedefs -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-stringop-overflow -Wno-ps\r\nabi -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -f\r\nno-math-errno -fno-trapping-math -Werror=format -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_VERSION=1.8.1, USE_CUDA=\r\nON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON,\r\n                                   \r\nTorchVision: 0.9.1+cu111  \r\nOpenCV: 4.5.3             \r\nMMCV: 1.3.8                \r\nMMCV Compiler: GCC 7.5\r\nMMCV CUDA Compiler: 11.1\r\nMMPose: 0.15.0+51b4b45\n", "before_files": [{"content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport warnings\n\nimport torch\nfrom mmcv.parallel import MMDataParallel, MMDistributedDataParallel\nfrom mmcv.runner import DistSamplerSeedHook, EpochBasedRunner, OptimizerHook\n\nfrom mmpose.core import DistEvalHook, EvalHook, build_optimizers\nfrom mmpose.core.distributed_wrapper import DistributedDataParallelWrapper\nfrom mmpose.datasets import build_dataloader, build_dataset\nfrom mmpose.utils import get_root_logger\n\ntry:\n    from mmcv.runner import Fp16OptimizerHook\nexcept ImportError:\n    warnings.warn(\n        'Fp16OptimizerHook from mmpose will be deprecated from '\n        'v0.15.0. Please install mmcv>=1.1.4', DeprecationWarning)\n    from mmpose.core import Fp16OptimizerHook\n\n\ndef train_model(model,\n                dataset,\n                cfg,\n                distributed=False,\n                validate=False,\n                timestamp=None,\n                meta=None):\n    \"\"\"Train model entry function.\n\n    Args:\n        model (nn.Module): The model to be trained.\n        dataset (Dataset): Train dataset.\n        cfg (dict): The config dict for training.\n        distributed (bool): Whether to use distributed training.\n            Default: False.\n        validate (bool): Whether to do evaluation. Default: False.\n        timestamp (str | None): Local time for runner. Default: None.\n        meta (dict | None): Meta dict to record some important information.\n            Default: None\n    \"\"\"\n    logger = get_root_logger(cfg.log_level)\n\n    # prepare data loaders\n    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]\n    dataloader_setting = dict(\n        samples_per_gpu=cfg.data.get('samples_per_gpu', {}),\n        workers_per_gpu=cfg.data.get('workers_per_gpu', {}),\n        # cfg.gpus will be ignored if distributed\n        num_gpus=len(cfg.gpu_ids),\n        dist=distributed,\n        seed=cfg.seed)\n    dataloader_setting = dict(dataloader_setting,\n                              **cfg.data.get('train_dataloader', {}))\n\n    data_loaders = [\n        build_dataloader(ds, **dataloader_setting) for ds in dataset\n    ]\n\n    # determine wether use adversarial training precess or not\n    use_adverserial_train = cfg.get('use_adversarial_train', False)\n\n    # put model on gpus\n    if distributed:\n        find_unused_parameters = cfg.get('find_unused_parameters', True)\n        # Sets the `find_unused_parameters` parameter in\n        # torch.nn.parallel.DistributedDataParallel\n\n        if use_adverserial_train:\n            # Use DistributedDataParallelWrapper for adversarial training\n            model = DistributedDataParallelWrapper(\n                model,\n                device_ids=[torch.cuda.current_device()],\n                broadcast_buffers=False,\n                find_unused_parameters=find_unused_parameters)\n        else:\n            model = MMDistributedDataParallel(\n                model.cuda(),\n                device_ids=[torch.cuda.current_device()],\n                broadcast_buffers=False,\n                find_unused_parameters=find_unused_parameters)\n    else:\n        model = MMDataParallel(\n            model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)\n\n    # build runner\n    optimizer = build_optimizers(model, cfg.optimizer)\n\n    runner = EpochBasedRunner(\n        model,\n        optimizer=optimizer,\n        work_dir=cfg.work_dir,\n        logger=logger,\n        meta=meta)\n    # an ugly workaround to make .log and .log.json filenames the same\n    runner.timestamp = timestamp\n\n    if use_adverserial_train:\n        # The optimizer step process is included in the train_step function\n        # of the model, so the runner should NOT include optimizer hook.\n        optimizer_config = None\n    else:\n        # fp16 setting\n        fp16_cfg = cfg.get('fp16', None)\n        if fp16_cfg is not None:\n            optimizer_config = Fp16OptimizerHook(\n                **cfg.optimizer_config, **fp16_cfg, distributed=distributed)\n        elif distributed and 'type' not in cfg.optimizer_config:\n            optimizer_config = OptimizerHook(**cfg.optimizer_config)\n        else:\n            optimizer_config = cfg.optimizer_config\n\n    # register hooks\n    runner.register_training_hooks(cfg.lr_config, optimizer_config,\n                                   cfg.checkpoint_config, cfg.log_config,\n                                   cfg.get('momentum_config', None))\n    if distributed:\n        runner.register_hook(DistSamplerSeedHook())\n\n    # register eval hooks\n    if validate:\n        eval_cfg = cfg.get('evaluation', {})\n        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))\n        dataloader_setting = dict(\n            samples_per_gpu=1,\n            workers_per_gpu=cfg.data.get('workers_per_gpu', {}),\n            # cfg.gpus will be ignored if distributed\n            num_gpus=len(cfg.gpu_ids),\n            dist=distributed,\n            drop_last=False,\n            shuffle=False)\n        dataloader_setting = dict(dataloader_setting,\n                                  **cfg.data.get('val_dataloader', {}))\n        val_dataloader = build_dataloader(val_dataset, **dataloader_setting)\n        eval_hook = DistEvalHook if distributed else EvalHook\n        runner.register_hook(eval_hook(val_dataloader, **eval_cfg))\n\n    if cfg.resume_from:\n        runner.resume(cfg.resume_from)\n    elif cfg.load_from:\n        runner.load_checkpoint(cfg.load_from)\n    runner.run(data_loaders, cfg.workflow, cfg.total_epochs)\n", "path": "mmpose/apis/train.py"}], "after_files": [{"content": "# Copyright (c) OpenMMLab. All rights reserved.\nimport warnings\n\nimport torch\nfrom mmcv.parallel import MMDataParallel, MMDistributedDataParallel\nfrom mmcv.runner import DistSamplerSeedHook, EpochBasedRunner, OptimizerHook\n\nfrom mmpose.core import DistEvalHook, EvalHook, build_optimizers\nfrom mmpose.core.distributed_wrapper import DistributedDataParallelWrapper\nfrom mmpose.datasets import build_dataloader, build_dataset\nfrom mmpose.utils import get_root_logger\n\ntry:\n    from mmcv.runner import Fp16OptimizerHook\nexcept ImportError:\n    warnings.warn(\n        'Fp16OptimizerHook from mmpose will be deprecated from '\n        'v0.15.0. Please install mmcv>=1.1.4', DeprecationWarning)\n    from mmpose.core import Fp16OptimizerHook\n\n\ndef train_model(model,\n                dataset,\n                cfg,\n                distributed=False,\n                validate=False,\n                timestamp=None,\n                meta=None):\n    \"\"\"Train model entry function.\n\n    Args:\n        model (nn.Module): The model to be trained.\n        dataset (Dataset): Train dataset.\n        cfg (dict): The config dict for training.\n        distributed (bool): Whether to use distributed training.\n            Default: False.\n        validate (bool): Whether to do evaluation. Default: False.\n        timestamp (str | None): Local time for runner. Default: None.\n        meta (dict | None): Meta dict to record some important information.\n            Default: None\n    \"\"\"\n    logger = get_root_logger(cfg.log_level)\n\n    # prepare data loaders\n    dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]\n    # step 1: give default values and override (if exist) from cfg.data\n    loader_cfg = {\n        **dict(\n            seed=cfg.get('seed'),\n            drop_last=False,\n            dist=distributed,\n            num_gpus=len(cfg.gpu_ids)),\n        **({} if torch.__version__ != 'parrots' else dict(\n               prefetch_num=2,\n               pin_memory=False,\n           )),\n        **dict((k, cfg.data[k]) for k in [\n                   'samples_per_gpu',\n                   'workers_per_gpu',\n                   'shuffle',\n                   'seed',\n                   'drop_last',\n                   'prefetch_num',\n                   'pin_memory',\n                   'persistent_workers',\n               ] if k in cfg.data)\n    }\n\n    # step 2: cfg.data.train_dataloader has highest priority\n    train_loader_cfg = dict(loader_cfg, **cfg.data.get('train_dataloader', {}))\n\n    data_loaders = [build_dataloader(ds, **train_loader_cfg) for ds in dataset]\n\n    # determine wether use adversarial training precess or not\n    use_adverserial_train = cfg.get('use_adversarial_train', False)\n\n    # put model on gpus\n    if distributed:\n        find_unused_parameters = cfg.get('find_unused_parameters', True)\n        # Sets the `find_unused_parameters` parameter in\n        # torch.nn.parallel.DistributedDataParallel\n\n        if use_adverserial_train:\n            # Use DistributedDataParallelWrapper for adversarial training\n            model = DistributedDataParallelWrapper(\n                model,\n                device_ids=[torch.cuda.current_device()],\n                broadcast_buffers=False,\n                find_unused_parameters=find_unused_parameters)\n        else:\n            model = MMDistributedDataParallel(\n                model.cuda(),\n                device_ids=[torch.cuda.current_device()],\n                broadcast_buffers=False,\n                find_unused_parameters=find_unused_parameters)\n    else:\n        model = MMDataParallel(\n            model.cuda(cfg.gpu_ids[0]), device_ids=cfg.gpu_ids)\n\n    # build runner\n    optimizer = build_optimizers(model, cfg.optimizer)\n\n    runner = EpochBasedRunner(\n        model,\n        optimizer=optimizer,\n        work_dir=cfg.work_dir,\n        logger=logger,\n        meta=meta)\n    # an ugly workaround to make .log and .log.json filenames the same\n    runner.timestamp = timestamp\n\n    if use_adverserial_train:\n        # The optimizer step process is included in the train_step function\n        # of the model, so the runner should NOT include optimizer hook.\n        optimizer_config = None\n    else:\n        # fp16 setting\n        fp16_cfg = cfg.get('fp16', None)\n        if fp16_cfg is not None:\n            optimizer_config = Fp16OptimizerHook(\n                **cfg.optimizer_config, **fp16_cfg, distributed=distributed)\n        elif distributed and 'type' not in cfg.optimizer_config:\n            optimizer_config = OptimizerHook(**cfg.optimizer_config)\n        else:\n            optimizer_config = cfg.optimizer_config\n\n    # register hooks\n    runner.register_training_hooks(cfg.lr_config, optimizer_config,\n                                   cfg.checkpoint_config, cfg.log_config,\n                                   cfg.get('momentum_config', None))\n    if distributed:\n        runner.register_hook(DistSamplerSeedHook())\n\n    # register eval hooks\n    if validate:\n        eval_cfg = cfg.get('evaluation', {})\n        val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))\n        dataloader_setting = dict(\n            samples_per_gpu=1,\n            workers_per_gpu=cfg.data.get('workers_per_gpu', {}),\n            # cfg.gpus will be ignored if distributed\n            num_gpus=len(cfg.gpu_ids),\n            dist=distributed,\n            drop_last=False,\n            shuffle=False)\n        dataloader_setting = dict(dataloader_setting,\n                                  **cfg.data.get('val_dataloader', {}))\n        val_dataloader = build_dataloader(val_dataset, **dataloader_setting)\n        eval_hook = DistEvalHook if distributed else EvalHook\n        runner.register_hook(eval_hook(val_dataloader, **eval_cfg))\n\n    if cfg.resume_from:\n        runner.resume(cfg.resume_from)\n    elif cfg.load_from:\n        runner.load_checkpoint(cfg.load_from)\n    runner.run(data_loaders, cfg.workflow, cfg.total_epochs)\n", "path": "mmpose/apis/train.py"}]}