support export jit trace and blade model for pose/keypoints & add pos…

…e model zoo (#294)

support export jit trace and blade model for pose/keypoints & add pose model zoo

README.md

@@ -232,6 +232,7 @@ Please refer to the following model zoo for more details.
 - [detection model zoo](docs/source/model_zoo_det.md)
 - [detection3d model zoo](docs/source/model_zoo_det3d.md)
 - [segmentation model zoo](docs/source/model_zoo_seg.md)
+- [pose model zoo](docs/source/model_zoo_pose.md)
 
 ## Data Hub
 

README_zh-CN.md

@@ -222,6 +222,7 @@ EasyCV是一个涵盖多个领域的基于Pytorch的计算机视觉工具箱，
 - [目标检测模型库](docs/source/model_zoo_det.md)
 - [3D目标检测模型库](docs/source/model_zoo_det3d.md)
 - [图像分割模型库](docs/source/model_zoo_seg.md)
+- [关键点模型库](docs/source/model_zoo_pose.md)
 
 ## 开源许可证
 

configs/pose/hrnet_w48_coco_256x192_udp.py

@@ -196,5 +196,14 @@
         data=dict(**data['val'], imgs_per_gpu=1),
         evaluators=[dict(type='CoCoPoseTopDownEvaluator', **evaluator_args)])
 ]
-export = dict(use_jit=False)
 checkpoint_sync_export = True
+export = dict(use_jit=False)
+# export = dict(
+#     type='blade',
+#     blade_config=dict(
+#         enable_fp16=True,
+#         fp16_fallback_op_ratio=0.0,
+#         customize_op_black_list=[
+#             'aten::select', 'aten::index', 'aten::slice', 'aten::view',
+#             'aten::upsample', 'aten::clamp', 'aten::clone', 'aten::add@-1'
+#         ]))

configs/pose/wholebody/hrnet_w48_coco_wholebody_384x288_dark_plus.py

@@ -185,6 +185,16 @@
         data=dict(**data['val'], imgs_per_gpu=1),
         evaluators=[dict(type='WholeBodyKeyPointEvaluator', **evaluator_args)])
 ]
-export = dict(use_jit=False)
 checkpoint_sync_export = True
 predict = dict(type='WholeBodyKeypointsPredictor', bbox_thr=0.8)
+
+export = dict(use_jit=False)
+# export = dict(
+#     type='blade',
+#     blade_config=dict(
+#         enable_fp16=True,
+#         fp16_fallback_op_ratio=0.0,
+#         customize_op_black_list=[
+#             'aten::select', 'aten::index', 'aten::slice', 'aten::view',
+#             'aten::upsample', 'aten::clamp', 'aten::clone', 'aten::add@-1'
+#         ]))

docs/source/model_zoo_pose.md

@@ -0,0 +1,33 @@
+# Pose/keypoints Model Zoo
+
+## Body
+
+**Topdown Heatmap + Hrnet + Udp on Coco**
+
+| Config                                                                                                                 | Params | AP    | AP50  | AP75  | AR    | AR50  | Download                                                                                                             |
+| ---------------------------------------------------------------------------------------------------------------------- | ------ | ----- | ----- | ----- | ----- | ----- | -------------------------------------------------------------------------------------------------------------------- |
+| [hrnet_w48_coco_256x192_udp](https://github.com/alibaba/EasyCV/blob/master/configs/pose/hrnet_w48_coco_256x192_udp.py) | 63.6M  | 0.781 | 0.936 | 0.849 | 0.801 | 0.941 | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/pose/top_down_hrnet/epoch_210.pth) |
+
+**Topdown Heatmap + Litehrnet on Coco**
+
+| Config                                                                                                               | Params | AP    | AP50  | AP75  | AR    | AR50  | Download                                                                                                                 |
+|:-------------------------------------------------------------------------------------------------------------------- | ------ | ----- | ----- | ----- | ----- | ----- | ------------------------------------------------------------------------------------------------------------------------ |
+| [litehrnet_30_coco_384x288](https://github.com/alibaba/EasyCV/blob/master/configs/pose/litehrnet_30_coco_384x288.py) | 1.76M  | 0.710 | 0.905 | 0.783 | 0.742 | 0.919 | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/pose/top_down_litehrnet/epoch_210.pth) |
+
+## WholeBody
+
+**Topdown Heatmap + Hrnet + Dark on Coco-Wholebody**
+
+| Config                                                                                                                                                           | Params | Body AP | Body AR | Foot AP | Foot AR | Face AP | Face AR | LeftHand AP | LeftHand AR | RightHand AP | RightHand AP | Whole AP | Whole AR | Download                                                                                                              |
+| ---------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------ | ------- | ------- | ------- | ------- | ------- | ------- | ----------- | ----------- | ------------ | ------------ | -------- | -------- | --------------------------------------------------------------------------------------------------------------------- |
+| [hrnet_w48_coco_wholebody_384x288_dark_plus](https://github.com/alibaba/EasyCV/blob/master/configs/pose/wholebody/hrnet_w48_coco_wholebody_384x288_dark_plus.py) | 63.6M  | 0.766   | 0.806   | 0.746   | 0.805   | 0.856   | 0.892   | 0.631       | 0.694       | 0.595        | 0.667        | 0.684    | 0.742    | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/pose/wholebody/hrnet/epoch_290.pth) |
+
+Note: `plus` means the  model is first pre-trained on original COCO dataset, and then fine-tuned on COCO-WholeBody dataset.
+
+## Hand
+
+**Topdown Heatmap + Hrnet + Dark + Coco + Wholebody on Coco_wholebody_hand**
+
+| Config                                                                                                                                                      | Params | PCK   | AUC   | EPE  | Download                                                                                                                 |
+| ----------------------------------------------------------------------------------------------------------------------------------------------------------- | ------ | ----- | ----- | ---- | ------------------------------------------------------------------------------------------------------------------------ |
+| [hrnet_w18_coco_wholebody_hand_256x256_dark](https://github.com/alibaba/EasyCV/blob/master/configs/pose/hand/hrnet_w18_coco_wholebody_hand_256x256_dark.py) | 9.6M   | 0.816 | 0.839 | 4.39 | [model](http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/pose/hand/hrnet/hrnet_w18_256x256.pth) |

easycv/apis/export.py

@@ -14,7 +14,7 @@
 from easycv.file import io
 from easycv.framework.errors import NotImplementedError, ValueError
 from easycv.models import (DINO, MOCO, SWAV, YOLOX, BEVFormer, Classification,
-                           MoBY, build_model)
+                           MoBY, TopDown, build_model)
 from easycv.utils.checkpoint import load_checkpoint
 from easycv.utils.misc import encode_str_to_tensor
 
@@ -66,6 +66,8 @@ def export(cfg, ckpt_path, filename, model=None, **kwargs):
         _export_yolox(model, cfg, filename, **kwargs)
     elif isinstance(model, BEVFormer):
         _export_bevformer(model, cfg, filename, **kwargs)
+    elif isinstance(model, TopDown):
+        _export_pose_topdown(model, cfg, filename, **kwargs)
     elif hasattr(cfg, 'export') and getattr(cfg.export, 'use_jit', False):
         export_jit_model(model, cfg, filename, **kwargs)
         return
@@ -632,6 +634,93 @@ def _get_blade_model():
         torch.jit.save(blade_model, ofile)
 
 
+def _export_pose_topdown(model, cfg, filename, fp16=False, dummy_inputs=None):
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model = copy.deepcopy(model)
+    model.eval()
+    model.to(device)
+
+    if hasattr(cfg, 'export') and getattr(cfg.export, 'type', 'raw') == 'raw':
+        return _export_common(model, cfg, filename)
+
+    def _dummy_inputs(cfg):
+        from easycv.datasets.pose.data_sources.top_down import DatasetInfo
+        from easycv.datasets.pose.data_sources.wholebody.wholebody_coco_source import WHOLEBODY_COCO_DATASET_INFO
+        from easycv.datasets.pose.data_sources.hand.coco_hand import COCO_WHOLEBODY_HAND_DATASET_INFO
+        from easycv.datasets.pose.data_sources.coco import COCO_DATASET_INFO
+
+        data_type = cfg.data.train.data_source.type
+        data_info_map = {
+            'WholeBodyCocoTopDownSource': WHOLEBODY_COCO_DATASET_INFO,
+            'PoseTopDownSourceCoco': COCO_DATASET_INFO,
+            'HandCocoPoseTopDownSource': COCO_WHOLEBODY_HAND_DATASET_INFO
+        }
+        dataset_info = DatasetInfo(data_info_map[data_type])
+        flip_pairs = dataset_info.flip_pairs
+        image_size = cfg.data_cfg.image_size
+        img = torch.rand([1, 3, image_size[1], image_size[0]]).to(device)
+        img_metas = [{
+            'image_id': torch.tensor(0),
+            'center': torch.tensor([426., 451.]),
+            'scale': torch.tensor([4., 5.]),
+            'rotation': torch.tensor(0),
+            'flip_pairs': torch.tensor(flip_pairs),
+            'bbox_id': torch.tensor(0)
+        }]
+        return img, img_metas
+
+    if dummy_inputs is None:
+        dummy_inputs = _dummy_inputs(cfg)
+
+    def _trace_model():
+        with torch.no_grad():
+            model.forward = model.forward_export
+            trace_model = torch.jit.trace(
+                model, copy.deepcopy(dummy_inputs), strict=False)
+        return trace_model
+
+    export_type = cfg.export.get('type')
+    if export_type in ['jit', 'blade']:
+        if fp16:
+            with torch.cuda.amp.autocast():
+                trace_model = _trace_model()
+        else:
+            trace_model = _trace_model()
+        torch.jit.save(trace_model, filename + '.jit')
+    else:
+        raise NotImplementedError(f'Not support export type {export_type}!')
+
+    if export_type == 'jit':
+        return
+
+    blade_config = cfg.export.get('blade_config')
+
+    from easycv.toolkit.blade import blade_env_assert, blade_optimize
+    assert blade_env_assert()
+
+    def _get_blade_model():
+        blade_model = blade_optimize(
+            speed_test_model=model,
+            model=trace_model,
+            inputs=copy.deepcopy(dummy_inputs),
+            blade_config=blade_config,
+            static_opt=False,
+            min_num_nodes=None,
+            check_inputs=False,
+            fp16=fp16)
+        return blade_model
+
+    # optimize model with blade
+    if fp16:
+        with torch.cuda.amp.autocast():
+            blade_model = _get_blade_model()
+    else:
+        blade_model = _get_blade_model()
+
+    with io.open(filename + '.blade', 'wb') as ofile:
+        torch.jit.save(blade_model, ofile)
+
+
 def replace_syncbn(backbone_cfg):
     if 'norm_cfg' in backbone_cfg.keys():
         if backbone_cfg['norm_cfg']['type'] == 'SyncBN':

easycv/core/post_processing/pose_transforms.py

@@ -4,6 +4,7 @@
 
 import cv2
 import numpy as np
+import torch
 
 
 def fliplr_joints(joints_3d, joints_3d_visible, img_width, flip_pairs):
@@ -125,20 +126,37 @@ def flip_back(output_flipped, flip_pairs, target_type='GaussianHeatmap'):
         'output_flipped should be [batch_size, num_keypoints, height, width]'
     shape_ori = output_flipped.shape
     channels = 1
-    if target_type.lower() == 'CombinedTarget'.lower():
-        channels = 3
-        output_flipped[:, 1::3, ...] = -output_flipped[:, 1::3, ...]
-    output_flipped = output_flipped.reshape(shape_ori[0], -1, channels,
-                                            shape_ori[2], shape_ori[3])
-    output_flipped_back = output_flipped.copy()
 
-    # Swap left-right parts
-    for left, right in flip_pairs:
-        output_flipped_back[:, left, ...] = output_flipped[:, right, ...]
-        output_flipped_back[:, right, ...] = output_flipped[:, left, ...]
-    output_flipped_back = output_flipped_back.reshape(shape_ori)
-    # Flip horizontally
-    output_flipped_back = output_flipped_back[..., ::-1]
+    if isinstance(output_flipped, torch.Tensor):
+        if target_type.lower() == 'CombinedTarget'.lower():
+            channels = 3
+            output_flipped[:, 1::3, ...] = -output_flipped[:, 1::3, ...]
+        output_flipped = output_flipped.reshape(shape_ori[0], -1, channels,
+                                                shape_ori[2], shape_ori[3])
+        output_flipped_back = output_flipped.clone()
+
+        # Swap left-right parts
+        for left, right in flip_pairs:
+            output_flipped_back[:, left, ...] = output_flipped[:, right, ...]
+            output_flipped_back[:, right, ...] = output_flipped[:, left, ...]
+        output_flipped_back = output_flipped_back.reshape(shape_ori)
+        # Flip horizontally
+        output_flipped_back = torch.flip(output_flipped_back, (-1, ))
+    else:
+        if target_type.lower() == 'CombinedTarget'.lower():
+            channels = 3
+            output_flipped[:, 1::3, ...] = -output_flipped[:, 1::3, ...]
+        output_flipped = output_flipped.reshape(shape_ori[0], -1, channels,
+                                                shape_ori[2], shape_ori[3])
+        output_flipped_back = output_flipped.copy()
+
+        # Swap left-right parts
+        for left, right in flip_pairs:
+            output_flipped_back[:, left, ...] = output_flipped[:, right, ...]
+            output_flipped_back[:, right, ...] = output_flipped[:, left, ...]
+        output_flipped_back = output_flipped_back.reshape(shape_ori)
+        # Flip horizontally
+        output_flipped_back = output_flipped_back[..., ::-1]
     return output_flipped_back
 
 

easycv/models/pose/heads/topdown_heatmap_base_head.py

@@ -10,6 +10,58 @@
 from easycv.framework.errors import ValueError
 
 
+def decode_heatmap(heatmaps, img_metas, test_cfg):
+    batch_size = len(img_metas)
+
+    if 'bbox_id' in img_metas[0]:
+        bbox_ids = []
+    else:
+        bbox_ids = None
+
+    c = np.zeros((batch_size, 2), dtype=np.float32)
+    s = np.zeros((batch_size, 2), dtype=np.float32)
+    image_ids = []
+    score = np.ones(batch_size)
+    for i in range(batch_size):
+        c[i, :] = img_metas[i]['center']
+        s[i, :] = img_metas[i]['scale']
+        image_ids.append(img_metas[i]['image_id'])
+
+        if 'bbox_score' in img_metas[i]:
+            score[i] = np.array(img_metas[i]['bbox_score']).reshape(-1)
+        if bbox_ids is not None:
+            bbox_ids.append(img_metas[i]['bbox_id'])
+
+    preds, maxvals = keypoints_from_heatmaps(
+        heatmaps,
+        c,
+        s,
+        unbiased=test_cfg.get('unbiased_decoding', False),
+        post_process=test_cfg.get('post_process', 'default'),
+        kernel=test_cfg.get('modulate_kernel', 11),
+        valid_radius_factor=test_cfg.get('valid_radius_factor', 0.0546875),
+        use_udp=test_cfg.get('use_udp', False),
+        target_type=test_cfg.get('target_type', 'GaussianHeatmap'))
+
+    all_preds = np.zeros((batch_size, preds.shape[1], 3), dtype=np.float32)
+    all_boxes = np.zeros((batch_size, 6), dtype=np.float32)
+    all_preds[:, :, 0:2] = preds[:, :, 0:2]
+    all_preds[:, :, 2:3] = maxvals
+    all_boxes[:, 0:2] = c[:, 0:2]
+    all_boxes[:, 2:4] = s[:, 0:2]
+    all_boxes[:, 4] = np.prod(s * 200.0, axis=1)
+    all_boxes[:, 5] = score
+
+    result = {}
+
+    result['preds'] = all_preds
+    result['boxes'] = all_boxes
+    result['image_ids'] = image_ids
+    result['bbox_ids'] = bbox_ids
+
+    return result
+
+
 class TopdownHeatmapBaseHead(nn.Module):
     """Base class for top-down heatmap heads.
 
@@ -53,56 +105,7 @@ def decode(self, img_metas, output, **kwargs):
                 - "bbox_score": score of bbox
             output (np.ndarray[N, K, H, W]): model predicted heatmaps.
         """
-        batch_size = len(img_metas)
-
-        if 'bbox_id' in img_metas[0]:
-            bbox_ids = []
-        else:
-            bbox_ids = None
-
-        c = np.zeros((batch_size, 2), dtype=np.float32)
-        s = np.zeros((batch_size, 2), dtype=np.float32)
-        image_ids = []
-        score = np.ones(batch_size)
-        for i in range(batch_size):
-            c[i, :] = img_metas[i]['center']
-            s[i, :] = img_metas[i]['scale']
-            image_ids.append(img_metas[i]['image_id'])
-
-            if 'bbox_score' in img_metas[i]:
-                score[i] = np.array(img_metas[i]['bbox_score']).reshape(-1)
-            if bbox_ids is not None:
-                bbox_ids.append(img_metas[i]['bbox_id'])
-
-        preds, maxvals = keypoints_from_heatmaps(
-            output,
-            c,
-            s,
-            unbiased=self.test_cfg.get('unbiased_decoding', False),
-            post_process=self.test_cfg.get('post_process', 'default'),
-            kernel=self.test_cfg.get('modulate_kernel', 11),
-            valid_radius_factor=self.test_cfg.get('valid_radius_factor',
-                                                  0.0546875),
-            use_udp=self.test_cfg.get('use_udp', False),
-            target_type=self.test_cfg.get('target_type', 'GaussianHeatmap'))
-
-        all_preds = np.zeros((batch_size, preds.shape[1], 3), dtype=np.float32)
-        all_boxes = np.zeros((batch_size, 6), dtype=np.float32)
-        all_preds[:, :, 0:2] = preds[:, :, 0:2]
-        all_preds[:, :, 2:3] = maxvals
-        all_boxes[:, 0:2] = c[:, 0:2]
-        all_boxes[:, 2:4] = s[:, 0:2]
-        all_boxes[:, 4] = np.prod(s * 200.0, axis=1)
-        all_boxes[:, 5] = score
-
-        result = {}
-
-        result['preds'] = all_preds
-        result['boxes'] = all_boxes
-        result['image_ids'] = image_ids
-        result['bbox_ids'] = bbox_ids
-
-        return result
+        return decode_heatmap(output, img_metas, self.test_cfg)
 
     @staticmethod
     def _get_deconv_cfg(deconv_kernel):

easycv/models/pose/heads/topdown_heatmap_simple_head.py

@@ -214,14 +214,17 @@ def inference_model(self, x, flip_pairs=None):
 
         if flip_pairs is not None:
             output_heatmap = flip_back(
-                output.detach().cpu().numpy(),
+                output.detach().cpu(),
                 flip_pairs,
                 target_type=self.target_type)
             # feature is not aligned, shift flipped heatmap for higher accuracy
             if self.test_cfg.get('shift_heatmap', False):
-                output_heatmap[:, :, :, 1:] = output_heatmap[:, :, :, :-1]
+                # remove inplace operation for blade, it will cause calculation errors
+                _tmp = output_heatmap[:, :, :, :-1]
+                output_heatmap[:, :, :, 1:] = _tmp
         else:
-            output_heatmap = output.detach().cpu().numpy()
+            output_heatmap = output.detach().cpu()
+
         return output_heatmap
 
     def _init_inputs(self, in_channels, in_index, input_transform):