Support fp16 training (open-mmlab#616)

* add fp16 configs * add auto_fp16 to detectors: TopDown, BottomUp and PoseLifter * Support fp16 training * switch to mmcv fp16 support * add fp16 decorators to top_down/bottom_up/pose_lifter * add config/ckpt of res50 fp16 * add config/ckpt of hrnet * fix bug
Spiideo · May 6, 2021 · 7cccdff · 7cccdff
1 parent b176078
commit 7cccdff
Show file tree

Hide file tree

Showing 19 changed files with 75 additions and 695 deletions.
diff --git a/configs/fp16/README.md b/configs/fp16/README.md
@@ -0,0 +1,14 @@
+# Mixed Precision Training
+
+## Introduction
+
+<!-- [OTHERS] -->
+
+```latex
+@article{micikevicius2017mixed,
+  title={Mixed precision training},
+  author={Micikevicius, Paulius and Narang, Sharan and Alben, Jonah and Diamos, Gregory and Elsen, Erich and Garcia, David and Ginsburg, Boris and Houston, Michael and Kuchaiev, Oleksii and Venkatesh, Ganesh and others},
+  journal={arXiv preprint arXiv:1710.03740},
+  year={2017}
+}
+```
diff --git a/configs/top_down/hrnet/README.md b/configs/top_down/hrnet/README.md
@@ -26,6 +26,9 @@
 | [pose_hrnet_w32](/configs/top_down/hrnet/coco/hrnet_w32_coco_384x288.py)  | 384x288 | 0.760 | 0.906 | 0.829 | 0.810 | 0.943 | [ckpt](https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w32_coco_384x288-d9f0d786_20200708.pth) | [log](https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w32_coco_384x288_20200708.log.json) |
 | [pose_hrnet_w48](/configs/top_down/hrnet/coco/hrnet_w48_coco_256x192.py)  | 256x192 | 0.756 | 0.907 | 0.825 | 0.806 | 0.942 | [ckpt](https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_256x192-b9e0b3ab_20200708.pth) | [log](https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_256x192_20200708.log.json) |
 | [pose_hrnet_w48](/configs/top_down/hrnet/coco/hrnet_w48_coco_384x288.py)  | 384x288 | 0.767 | 0.910 | 0.831 | 0.816 | 0.946 | [ckpt](https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_384x288-314c8528_20200708.pth) | [log](https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w48_coco_384x288_20200708.log.json) |
+| [pose_hrnet_w32_fp16<sup>1</sup>](/configs/top_down/hrnet/coco/hrnet_w32_coco_256x192_fp16_dynamic.py)  | 256x192 | 0.746 | 0.905 | 0.88 | 0.800 | 0.943 | [ckpt](hrnet_w32_coco_256x192_fp16_dynamic-290efc2e_20210430.pth) | [log](https://download.openmmlab.com/mmpose/top_down/hrnet/hrnet_w32_coco_256x192_fp16_dynamic_20210430.log.json) |
+
+<sup>1</sup> Please refer to [fp16/README.md](/configs/fp16/README.md) for the method we use for mixed precision training.
 
 #### Results on AIC val set with ground-truth bounding boxes
 

diff --git a/configs/top_down/hrnet/coco/hrnet_w32_coco_256x192_fp16_dynamic.py b/configs/top_down/hrnet/coco/hrnet_w32_coco_256x192_fp16_dynamic.py
@@ -0,0 +1,4 @@
+_base_ = ['./hrnet_w32_coco_256x192.py']
+
+# fp16 settings
+fp16 = dict(loss_scale='dynamic')
diff --git a/configs/top_down/resnet/README.md b/configs/top_down/resnet/README.md
@@ -28,6 +28,9 @@
 | [pose_resnet_101](/configs/top_down/resnet/coco/res101_coco_384x288.py) | 384x288 | 0.748 | 0.905 | 0.817 | 0.798 | 0.940 | [ckpt](https://download.openmmlab.com/mmpose/top_down/resnet/res101_coco_384x288-8c71bdc9_20200709.pth) | [log](https://download.openmmlab.com/mmpose/top_down/resnet/res101_coco_384x288_20200709.log.json) |
 | [pose_resnet_152](/configs/top_down/resnet/coco/res152_coco_256x192.py) | 256x192 | 0.735 | 0.905 | 0.812 | 0.790 | 0.943 | [ckpt](https://download.openmmlab.com/mmpose/top_down/resnet/res152_coco_256x192-f6e307c2_20200709.pth) | [log](https://download.openmmlab.com/mmpose/top_down/resnet/res152_coco_256x192_20200709.log.json) |
 | [pose_resnet_152](/configs/top_down/resnet/coco/res152_coco_384x288.py) | 384x288 | 0.750 | 0.908 | 0.821 | 0.800 | 0.942 | [ckpt](https://download.openmmlab.com/mmpose/top_down/resnet/res152_coco_384x288-3860d4c9_20200709.pth) | [log](https://download.openmmlab.com/mmpose/top_down/resnet/res152_coco_384x288_20200709.log.json) |
+| [pose_resnet_50_fp16<sup>1</sup>](/configs/top_down/resnet/coco/res50_coco_256x192_fp16_dynamic.py)  | 256x192 | 0.717 | 0.898 | 0.793 | 0.772 | 0.936 | [ckpt](https://download.openmmlab.com/mmpose/top_down/resnet/res50_coco_256x192_fp16_dynamic-6edb79f3_20210430.pth) | [log](https://download.openmmlab.com/mmpose/top_down/resnet/res50_coco_256x192_fp16_dynamic_20210430.log.json) |
+
+<sup>1</sup> Please refer to [fp16/README.md](/configs/fp16/README.md) for the method we use for mixed precision training.
 
 #### Results on OCHuman test dataset with ground-truth bounding boxes
 

diff --git a/configs/top_down/resnet/coco/res50_coco_256x192_fp16_dynamic.py b/configs/top_down/resnet/coco/res50_coco_256x192_fp16_dynamic.py
@@ -0,0 +1,4 @@
+_base_ = ['./res50_coco_256x192.py']
+
+# fp16 settings
+fp16 = dict(loss_scale='dynamic')
diff --git a/mmpose/apis/train.py b/mmpose/apis/train.py
@@ -1,23 +1,19 @@
 import warnings
 
-import mmcv
 import torch
 from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
 from mmcv.runner import DistSamplerSeedHook, EpochBasedRunner, OptimizerHook
-from packaging import version
 
 from mmpose.core import DistEvalHook, EvalHook, build_optimizers
 from mmpose.core.distributed_wrapper import DistributedDataParallelWrapper
 from mmpose.datasets import build_dataloader, build_dataset
 from mmpose.utils import get_root_logger
 
-if version.parse(mmcv.__version__) >= version.parse('1.3.1'):
+try:
     from mmcv.runner import Fp16OptimizerHook
-else:
-    warnings.warn(
-        'Fp16OptimizerHook in mmpose will be deprecated in the next'
-        'release and replaced by Fp16OptimizerHook provided by mmcv (>=1.3.1)'
-        'Please upgrade your mmcv version, if needed.', DeprecationWarning)
+except ImportError:
+    warnings.warn('Fp16OptimizerHook from mmpose will be deprecated from '
+                  'v0.15.0. Please install mmcv>=1.1.4')
     from mmpose.core import Fp16OptimizerHook
 
 

diff --git a/mmpose/core/__init__.py b/mmpose/core/__init__.py
@@ -1,6 +1,5 @@
 from .camera import *  # noqa: F401, F403
 from .evaluation import *  # noqa: F401, F403
-from .fp16 import *  # noqa: F401, F403
 from .optimizer import *  # noqa: F401, F403
 from .post_processing import *  # noqa: F401, F403
 from .utils import *  # noqa: F401, F403
diff --git a/mmpose/core/fp16/__init__.py b/mmpose/core/fp16/__init__.py
diff --git a/mmpose/core/fp16/decorators.py b/mmpose/core/fp16/decorators.py