[Enhancement] Support litehrnet for ncnn (open-mmlab#316)

* support for litehrnet

* support ncnn vulkan

* update supported models

* add docstring

* add one ut and fix wrapper

* add test_shuffleunit

* add test_stem_forward

* add last ut

* fix flake8

* fix yapf

* fix lint

* fix yapf

* fix comments

* unified chunk

* mv adaptive_avg_pool

* fix lint

* move adaptive_avg_pool_ncnn

* fix lint

* symbolic rewriter of adaptive_avg_pool2d

* fix lint

* fix yapf

* fix ci

docs/en/codebases/mmpose.md

@@ -12,7 +12,7 @@ Please refer to [official installation guide](https://mmpose.readthedocs.io/en/l
 |:----------|:--------------|:------------:|:--------:|:----:|:-----:|:--------:|:-------------------------------------------------------------------------------------------:|
 | HRNet     | PoseDetection |      Y       |    Y     |  Y   |   N   |    Y     |   [config](https://mmpose.readthedocs.io/en/latest/papers/backbones.html#hrnet-cvpr-2019)   |
 | MSPN      | PoseDetection |      Y       |    Y     |  Y   |   N   |    Y     |   [config](https://mmpose.readthedocs.io/en/latest/papers/backbones.html#mspn-arxiv-2019)   |
-| LiteHRNet | PoseDetection |      Y       |    Y     |  N   |   N   |    Y     | [config](https://mmpose.readthedocs.io/en/latest/papers/backbones.html#litehrnet-cvpr-2021) |
+| LiteHRNet | PoseDetection |      Y       |    Y     |  Y   |   N   |    Y     | [config](https://mmpose.readthedocs.io/en/latest/papers/backbones.html#litehrnet-cvpr-2021) |
 
 ### Example
 

docs/en/supported_models.md

@@ -63,7 +63,7 @@ The table below lists the models that are guaranteed to be exportable to other b
 | SAR                       | MMOCR            |      N      |      Y      |     N    |   N   |   N   |     N    | [config](https://github.com/open-mmlab/mmocr/tree/main/configs/textrecog/sar)                  |
 | HRNet                     | MMPose           |      N      |      Y      |     Y    |   Y   |   N   |     Y    |    [config](https://mmpose.readthedocs.io/en/latest/papers/backbones.html#hrnet-cvpr-2019)     |
 | MSPN                      | MMPose           |      N      |      Y      |     Y    |   Y   |   N   |     Y    |    [config](https://mmpose.readthedocs.io/en/latest/papers/backbones.html#mspn-arxiv-2019)     |
-| LiteHRNet                 | MMPose           |      N      |      Y      |     Y    |   N   |   N   |     Y    |  [config](https://mmpose.readthedocs.io/en/latest/papers/backbones.html#litehrnet-cvpr-2021)   |
+| LiteHRNet                 | MMPose           |      N      |      Y      |     Y    |   Y   |   N   |     Y    |  [config](https://mmpose.readthedocs.io/en/latest/papers/backbones.html#litehrnet-cvpr-2021)   |
 | PointPillars              | MMDetection3d    |      ?      |      Y      |     Y    |   N   |   N   |     Y    | [config](https://github.com/open-mmlab/mmdetection3d/blob/master/configs/pointpillars)         |
 | CenterPoint (pillar)      | MMDetection3d    |      ?      |      Y      |     Y    |   N   |   N   |     Y    | [config](https://github.com/open-mmlab/mmdetection3d/blob/master/configs/centerpoint)          |
 

mmdeploy/codebase/base/backend_model.py

@@ -74,7 +74,14 @@ def _build_wrapper(backend: Backend,
                 output_names=output_names)
         elif backend == Backend.NCNN:
             from mmdeploy.backend.ncnn import NCNNWrapper
-            use_vulkan = get_backend_config('use_vulkan', False)
+
+            # For unittest deploy_config will not pass into _build_wrapper
+            # function.
+            if deploy_cfg:
+                backend_config = get_backend_config(deploy_cfg)
+                use_vulkan = backend_config.get('use_vulkan', False)
+            else:
+                use_vulkan = False
             return NCNNWrapper(
                 param_file=backend_files[0],
                 bin_file=backend_files[1],

mmdeploy/pytorch/functions/__init__.py

@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+from .chunk import chunk__ncnn
 from .getattribute import tensor__getattribute__ncnn
 from .group_norm import group_norm__ncnn
 from .interpolate import interpolate__ncnn, interpolate__tensorrt
@@ -10,5 +11,5 @@
 __all__ = [
     'tensor__getattribute__ncnn', 'group_norm__ncnn', 'interpolate__ncnn',
     'interpolate__tensorrt', 'linear__ncnn', 'tensor__repeat__tensorrt',
-    'tensor__size__ncnn', 'topk__dynamic', 'topk__tensorrt'
+    'tensor__size__ncnn', 'topk__dynamic', 'topk__tensorrt', 'chunk__ncnn'
 ]

mmdeploy/pytorch/functions/chunk.py

@@ -0,0 +1,33 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import torch
+
+from mmdeploy.core import FUNCTION_REWRITER
+
+
+@FUNCTION_REWRITER.register_rewriter(
+    func_name='torch.Tensor.chunk', backend='ncnn')
+def chunk__ncnn(ctx, self, num_chunks: int, dim: int = 0) -> torch.Tensor:
+    """Rewrite `chunk` for NCNN backend.
+
+    Chunk in ncnn are not supported, so it should be rewritten.
+    """
+    dim_len = self.shape[dim]
+    # int ceil.
+    step = dim_len // num_chunks
+    if dim_len % num_chunks > 0:
+        step += 1
+    index_list = []
+    index = 0
+    while index < dim_len:
+        index_list.append(index)
+        index += step
+    index_list.append(dim_len)
+    output = [
+        self.index_select(
+            dim,
+            torch.tensor([j for j in range(index_list[i], index_list[i + 1])],
+                         dtype=torch.int64))
+        for i in range(len(index_list) - 1)
+    ]
+
+    return output

mmdeploy/pytorch/ops/__init__.py

@@ -1,6 +1,7 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from .adaptive_avg_pool import (adaptive_avg_pool1d__default,
                                 adaptive_avg_pool2d__default,
+                                adaptive_avg_pool2d__ncnn,
                                 adaptive_avg_pool3d__default)
 from .gelu import gelu__ncnn
 from .grid_sampler import grid_sampler__default
@@ -14,5 +15,6 @@
     'adaptive_avg_pool1d__default', 'adaptive_avg_pool2d__default',
     'adaptive_avg_pool3d__default', 'grid_sampler__default',
     'hardsigmoid__default', 'instance_norm__tensorrt', 'generic_rnn__ncnn',
-    'squeeze__default', 'gelu__ncnn', 'layer_norm__ncnn'
+    'squeeze__default', 'adaptive_avg_pool2d__ncnn', 'gelu__ncnn',
+    'layer_norm__ncnn'
 ]

mmdeploy/pytorch/ops/adaptive_avg_pool.py

@@ -78,3 +78,13 @@ def adaptive_avg_pool3d__default(ctx, *args):
     Align symbolic of adaptive_pool between different torch version.
     """
     return adaptive_avg_pool3d(*args)
+
+
+@SYMBOLIC_REWRITER.register_symbolic(
+    'adaptive_avg_pool2d', is_pytorch=True, backend='ncnn')
+def adaptive_avg_pool2d__ncnn(ctx, g, x, output_size):
+    """Register ncnn symbolic function for `adaptive_avg_pool2d`.
+
+    Align symbolic of adaptive_avg_pool2d in ncnn.
+    """
+    return g.op('mmdeploy::adaptive_avg_pool2d', x, output_size)

tests/test_codebase/test_mmpose/test_mmpose_models.py

@@ -170,19 +170,29 @@ def forward(self, x):
     return model
 
 
-@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME])
+@pytest.mark.parametrize('backend_type', [Backend.ONNXRUNTIME, Backend.NCNN])
 def test_cross_resolution_weighting_forward(backend_type: Backend):
     check_backend(backend_type, True)
     model = get_cross_resolution_weighting_model()
     model.cpu().eval()
     imgs = torch.rand(1, 16, 16, 16)
 
-    deploy_cfg = mmcv.Config(
-        dict(
-            backend_config=dict(type=backend_type.value),
-            onnx_config=dict(input_shape=None, output_names=['output']),
-            codebase_config=dict(
-                type=Codebase.MMPOSE.value, task=Task.POSE_DETECTION.value)))
+    if backend_type == Backend.NCNN:
+        deploy_cfg = mmcv.Config(
+            dict(
+                backend_config=dict(type=backend_type.value, use_vulkan=False),
+                onnx_config=dict(input_shape=None, output_names=['output']),
+                codebase_config=dict(
+                    type=Codebase.MMPOSE.value,
+                    task=Task.POSE_DETECTION.value)))
+    else:
+        deploy_cfg = mmcv.Config(
+            dict(
+                backend_config=dict(type=backend_type.value),
+                onnx_config=dict(input_shape=None, output_names=['output']),
+                codebase_config=dict(
+                    type=Codebase.MMPOSE.value,
+                    task=Task.POSE_DETECTION.value)))
     rewrite_inputs = {'x': imgs}
     model_outputs = model.forward(imgs)
     wrapped_model = WrapModel(model, 'forward')

tests/test_pytorch/test_pytorch_functions.py

@@ -12,7 +12,7 @@
 deploy_cfg_ncnn = mmcv.Config(
     dict(
         onnx_config=dict(input_shape=None),
-        backend_config=dict(type='ncnn', model_inputs=None),
+        backend_config=dict(type='ncnn', model_inputs=None, use_vulkan=False),
         codebase_config=dict(type='mmdet', task='ObjectDetection')))
 
 
@@ -76,6 +76,28 @@ def group_norm_caller(input):
     assert np.allclose(model_output, rewrite_output[0], rtol=1e-03, atol=1e-05)
 
 
+@backend_checker(Backend.NCNN)
+def test_chunk_ncnn():
+    input = torch.rand(1, 16, 16, 16)
+
+    model_output = input.chunk(2, dim=1)
+
+    def chunk_caller(input):
+        return input.chunk(2, dim=1)
+
+    wrapped_func = WrapFunction(chunk_caller)
+    rewrite_output, _ = get_rewrite_outputs(
+        wrapped_func,
+        model_inputs={'input': input},
+        deploy_cfg=deploy_cfg_ncnn,
+        run_with_backend=True)
+
+    assert len(model_output) == len(rewrite_output)
+    for i in range(len(model_output)):
+        assert np.allclose(
+            model_output[i], rewrite_output[i], rtol=1e-03, atol=1e-05)
+
+
 @backend_checker(Backend.NCNN)
 def test_interpolate_static():
     input = torch.rand([1, 2, 2, 2])

tests/test_pytorch/test_pytorch_ops.py

@@ -11,7 +11,7 @@
 onnx_file = tempfile.NamedTemporaryFile(suffix='onnx').name
 
 
-@pytest.fixture(autouse=True, scope='module')
+@pytest.fixture(autouse=False, scope='function')
 def prepare_symbolics():
     context = RewriterContext(
         Config({'backend_config': {
@@ -24,6 +24,19 @@ def prepare_symbolics():
     context.exit()
 
 
+@pytest.fixture(autouse=False, scope='function')
+def prepare_symbolics_ncnn():
+    context = RewriterContext(
+        Config({'backend_config': {
+            'type': 'ncnn'
+        }}), 'ncnn', opset=11)
+    context.enter()
+
+    yield
+
+    context.exit()
+
+
 class OpModel(torch.nn.Module):
 
     def __init__(self, func, *args):
@@ -42,6 +55,7 @@ def get_model_onnx_nodes(model, x, onnx_file=onnx_file):
     return nodes
 
 
+@pytest.mark.usefixtures('prepare_symbolics')
 class TestAdaptivePool:
 
     def test_adaptive_pool_1d_global(self):
@@ -83,6 +97,17 @@ def test_adaptive_pool_3d(self):
         assert nodes[0].op_type == 'AveragePool'
 
 
+@pytest.mark.usefixtures('prepare_symbolics_ncnn')
+def test_adaptive_pool_2d_ncnn():
+    x = torch.rand(2, 2, 2)
+    model = OpModel(torch.nn.functional.adaptive_avg_pool2d,
+                    torch.tensor([2, 2], dtype=torch.int64)).eval()
+    nodes = get_model_onnx_nodes(model, x)
+    assert nodes[1].op_type == 'adaptive_avg_pool2d'
+    assert nodes[1].domain == 'mmdeploy'
+
+
+@pytest.mark.usefixtures('prepare_symbolics')
 def test_grid_sampler():
     x = torch.rand(1, 1, 2, 2)
     flow = torch.zeros([1, 2, 2, 2])
@@ -92,6 +117,7 @@ def test_grid_sampler():
     assert nodes[1].domain == 'mmdeploy'
 
 
+@pytest.mark.usefixtures('prepare_symbolics')
 def test_instance_norm():
     x = torch.rand(1, 2, 2, 2)
     model = OpModel(torch.group_norm, 1, torch.rand([2]), torch.rand([2]),
@@ -101,6 +127,7 @@ def test_instance_norm():
     assert nodes[4].domain == 'mmdeploy'
 
 
+@pytest.mark.usefixtures('prepare_symbolics')
 class TestSqueeze:
 
     def test_squeeze_default(self):
@@ -118,6 +145,7 @@ def test_squeeze(self):
         assert nodes[0].op_type == 'Squeeze'
 
 
+@pytest.mark.usefixtures('prepare_symbolics')
 def test_hardsigmoid():
     x = torch.rand(1, 2, 3, 4)
     model = torch.nn.Hardsigmoid().eval()