[Bug fix] fixed fp16 inference (open-mmlab#497)

xvjiarui · web-flow · commit 89f8a6a41003 · 2021-04-24T09:57:33.000-07:00
* fixed fp16

* update fps
diff --git a/configs/fp16/README.md b/configs/fp16/README.md
@@ -19,7 +19,7 @@
 
 | Method     | Backbone | Crop Size | Lr schd | Mem (GB) | Inf time (fps) |  mIoU | mIoU(ms+flip) | config                                                                                                                                | download                                                                                                                                                                                                                                                                                                                                                                                 |
 | ---------- | -------- | --------- | ------: | -------- | -------------- | ----: | ------------: | ------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| FCN        | R-101-D8 | 512x1024  |   80000 | 5.50     | 2.66           | 76.80 |             - | [config](https://github.com/open-mmlab/mmsegmentation/blob/master/configs/fp16/fcn_r101-d8_512x1024_80k_fp16_cityscapes.py)           | [model](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/fcn_r101-d8_512x1024_80k_fp16_cityscapes/fcn_r101-d8_512x1024_80k_fp16_cityscapes-50245227.pth) &#124; [log](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/fcn_r101-d8_512x1024_80k_fp16_cityscapes/fcn_r101-d8_512x1024_80k_fp16_cityscapes_20200717_230921.log.json)                                         |
-| PSPNet     | R-101-D8 | 512x1024  |   80000 | 5.47     | 2.68           | 79.46 |             - | [config](https://github.com/open-mmlab/mmsegmentation/blob/master/configs/fp16/pspnet_r101-d8_512x1024_80k_fp16_cityscapes.py)        | [model](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/pspnet_r101-d8_512x1024_80k_fp16_cityscapes/pspnet_r101-d8_512x1024_80k_fp16_cityscapes-ade37931.pth) &#124; [log](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/pspnet_r101-d8_512x1024_80k_fp16_cityscapes/pspnet_r101-d8_512x1024_80k_fp16_cityscapes_20200717_230919.log.json)                             |
-| DeepLabV3  | R-101-D8 | 512x1024  |   80000 | 5.91     | 1.93           | 80.48 |             - | [config](https://github.com/open-mmlab/mmsegmentation/blob/master/configs/fp16/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes.py)     | [model](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes-bc86dc84.pth) &#124; [log](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes_20200717_230920.log.json)                 |
-| DeepLabV3+ | R-101-D8 | 512x1024  |   80000 | 6.46     | 2.60           | 80.46 |             - | [config](https://github.com/open-mmlab/mmsegmentation/blob/master/configs/fp16/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes.py) | [model](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes-cc58bc8d.pth) &#124; [log](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes_20200717_230920.log.json) |
+| FCN        | R-101-D8 | 512x1024  |   80000 | 5.37     | 8.64           | 76.80 |             - | [config](https://github.com/open-mmlab/mmsegmentation/blob/master/configs/fp16/fcn_r101-d8_512x1024_80k_fp16_cityscapes.py)           | [model](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/fcn_r101-d8_512x1024_80k_fp16_cityscapes/fcn_r101-d8_512x1024_80k_fp16_cityscapes-50245227.pth) &#124; [log](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/fcn_r101-d8_512x1024_80k_fp16_cityscapes/fcn_r101-d8_512x1024_80k_fp16_cityscapes_20200717_230921.log.json)                                         |
+| PSPNet     | R-101-D8 | 512x1024  |   80000 | 5.34     | 8.77           | 79.46 |             - | [config](https://github.com/open-mmlab/mmsegmentation/blob/master/configs/fp16/pspnet_r101-d8_512x1024_80k_fp16_cityscapes.py)        | [model](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/pspnet_r101-d8_512x1024_80k_fp16_cityscapes/pspnet_r101-d8_512x1024_80k_fp16_cityscapes-ade37931.pth) &#124; [log](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/pspnet_r101-d8_512x1024_80k_fp16_cityscapes/pspnet_r101-d8_512x1024_80k_fp16_cityscapes_20200717_230919.log.json)                             |
+| DeepLabV3  | R-101-D8 | 512x1024  |   80000 | 5.75     | 3.86           | 80.48 |             - | [config](https://github.com/open-mmlab/mmsegmentation/blob/master/configs/fp16/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes.py)     | [model](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes-bc86dc84.pth) &#124; [log](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes_20200717_230920.log.json)                 |
+| DeepLabV3+ | R-101-D8 | 512x1024  |   80000 | 6.35     | 7.87           | 80.46 |             - | [config](https://github.com/open-mmlab/mmsegmentation/blob/master/configs/fp16/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes.py) | [model](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes-cc58bc8d.pth) &#124; [log](https://download.openmmlab.com/mmsegmentation/v0.5/fp16/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes_20200717_230920.log.json) |
diff --git a/configs/fp16/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes.py b/configs/fp16/deeplabv3_r101-d8_512x1024_80k_fp16_cityscapes.py
@@ -1,3 +1,5 @@
 _base_ = '../deeplabv3/deeplabv3_r101-d8_512x1024_80k_cityscapes.py'
 # fp16 settings
 optimizer_config = dict(type='Fp16OptimizerHook', loss_scale=512.)
+# fp16 placeholder
+fp16 = dict()
diff --git a/configs/fp16/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes.py b/configs/fp16/deeplabv3plus_r101-d8_512x1024_80k_fp16_cityscapes.py
@@ -1,3 +1,5 @@
 _base_ = '../deeplabv3plus/deeplabv3plus_r101-d8_512x1024_80k_cityscapes.py'
 # fp16 settings
 optimizer_config = dict(type='Fp16OptimizerHook', loss_scale=512.)
+# fp16 placeholder
+fp16 = dict()
diff --git a/configs/fp16/fcn_r101-d8_512x1024_80k_fp16_cityscapes.py b/configs/fp16/fcn_r101-d8_512x1024_80k_fp16_cityscapes.py
@@ -1,3 +1,5 @@
 _base_ = '../fcn/fcn_r101-d8_512x1024_80k_cityscapes.py'
 # fp16 settings
 optimizer_config = dict(type='Fp16OptimizerHook', loss_scale=512.)
+# fp16 placeholder
+fp16 = dict()
diff --git a/configs/fp16/pspnet_r101-d8_512x1024_80k_fp16_cityscapes.py b/configs/fp16/pspnet_r101-d8_512x1024_80k_fp16_cityscapes.py
@@ -1,3 +1,5 @@
 _base_ = '../pspnet/pspnet_r101-d8_512x1024_80k_cityscapes.py'
 # fp16 settings
 optimizer_config = dict(type='Fp16OptimizerHook', loss_scale=512.)
+# fp16 placeholder
+fp16 = dict()
diff --git a/tools/benchmark.py b/tools/benchmark.py
@@ -4,7 +4,7 @@
 import torch
 from mmcv import Config
 from mmcv.parallel import MMDataParallel
-from mmcv.runner import load_checkpoint
+from mmcv.runner import load_checkpoint, wrap_fp16_model
 
 from mmseg.datasets import build_dataloader, build_dataset
 from mmseg.models import build_segmentor
@@ -42,6 +42,9 @@ def main():
     # build the model and load checkpoint
     cfg.model.train_cfg = None
     model = build_segmentor(cfg.model, test_cfg=cfg.get('test_cfg'))
+    fp16_cfg = cfg.get('fp16', None)
+    if fp16_cfg is not None:
+        wrap_fp16_model(model)
     load_checkpoint(model, args.checkpoint, map_location='cpu')
 
     model = MMDataParallel(model, device_ids=[0])
diff --git a/tools/test.py b/tools/test.py
@@ -4,7 +4,8 @@
 import mmcv
 import torch
 from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
-from mmcv.runner import get_dist_info, init_dist, load_checkpoint
+from mmcv.runner import (get_dist_info, init_dist, load_checkpoint,
+                         wrap_fp16_model)
 from mmcv.utils import DictAction
 
 from mmseg.apis import multi_gpu_test, single_gpu_test
@@ -117,6 +118,9 @@ def main():
     # build the model and load checkpoint
     cfg.model.train_cfg = None
     model = build_segmentor(cfg.model, test_cfg=cfg.get('test_cfg'))
+    fp16_cfg = cfg.get('fp16', None)
+    if fp16_cfg is not None:
+        wrap_fp16_model(model)
     checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
     model.CLASSES = checkpoint['meta']['CLASSES']
     model.PALETTE = checkpoint['meta']['PALETTE']
