Improve ncnn patch embed (#592)

* improvement(PatchEmbed): use reshape instead of flatten

* improvement(init_plugins.py): get onnx2ncnn path

* fix(CI): build error

* fix(CI): revert get onnx2ncnn path

README_zh-CN.md

@@ -19,7 +19,7 @@
   <div>&nbsp;</div>
 </div>
 
-[![docs](https://img.shields.io/badge/docs-latest-blue)](https://mmdeploy.readthedocs.io/en/latest/)
+[![docs](https://img.shields.io/badge/docs-latest-blue)](https://mmdeploy.readthedocs.io/zh_CN/latest/)
 [![badge](https://github.com/open-mmlab/mmdeploy/workflows/build/badge.svg)](https://github.com/open-mmlab/mmdeploy/actions)
 [![codecov](https://codecov.io/gh/open-mmlab/mmdeploy/branch/master/graph/badge.svg)](https://codecov.io/gh/open-mmlab/mmdeploy)
 [![license](https://img.shields.io/github/license/open-mmlab/mmdeploy.svg)](https://github.com/open-mmlab/mmdeploy/blob/master/LICENSE)

mmdeploy/mmcv/ops/__init__.py

@@ -5,8 +5,10 @@ from .nms import *  # noqa: F401,F403
 from .nms_rotated import *  # noqa: F401,F403
 from .roi_align import roi_align_default
 from .roi_align_rotated import roi_align_rotated_default
+from .transformer import patch_embed__forward__ncnn
 
 __all__ = [
     'roi_align_default', 'modulated_deform_conv_default',
-    'deform_conv_openvino', 'roi_align_rotated_default'
+    'deform_conv_openvino', 'roi_align_rotated_default',
+    'patch_embed__forward__ncnn'
 ]

mmdeploy/mmcv/ops/transformer.py

@@ -0,0 +1,32 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+from mmdeploy.core import FUNCTION_REWRITER
+from mmdeploy.utils import Backend
+
+
+@FUNCTION_REWRITER.register_rewriter(
+    func_name='mmcv.cnn.bricks.transformer.PatchEmbed.forward',
+    backend=Backend.NCNN.value)
+def patch_embed__forward__ncnn(ctx, self, x):
+    """Rewrite `forward` of PatchEmbed for ncnn backend.
+
+    Args:
+        x (Tensor): Has shape (B, C, H, W). In most case, C is 3.
+
+    Returns:
+        tuple: Contains merged results and its spatial shape.
+
+        - x (Tensor): Has shape (B, out_h * out_w, embed_dims)
+        - out_size (tuple[int]): Spatial shape of x, arrange as
+            (out_h, out_w).
+    """
+
+    if self.adaptive_padding:
+        x = self.adaptive_padding(x)
+
+    x = self.projection(x)
+    x_shape = x.shape
+    out_size = (x_shape[2], x_shape[3])
+    x = x.reshape((x_shape[0], x_shape[1], -1)).transpose(1, 2)
+    if self.norm is not None:
+        x = self.norm(x)
+    return x, out_size

mmdeploy/utils/utils.py

@@ -115,7 +115,7 @@ def get_file_path(prefix, candidates) -> str:
 
     Args:
         prefix (str): Prefix of the paths.
-        cancidates (str): Candidate paths
+        candidates (str): Candidate paths
     Returns:
         str: file path or '' if not found
     """

tests/test_mmcv/test_mmcv_ops.py

@@ -93,3 +93,24 @@ def test_deform_conv_openvino():
     model = onnx.load(onnx_file_path)
     assert model.graph.node[1].op_type == 'DeformableConv2D'
     assert model.graph.node[1].domain == 'org.openvinotoolkit'
+
+
+def test_patch_embed_ncnn():
+    check_backend(Backend.NCNN)
+
+    from mmcv.cnn.bricks.transformer import PatchEmbed
+
+    input = torch.ones((1, 3, 384, 384))
+    patch_cfg = {
+        'in_channels': 3,
+        'input_size': 384,
+        'embed_dims': 768,
+        'conv_type': 'Conv2d',
+        'kernel_size': 32,
+        'stride': 32
+    }
+    wrapped_model = PatchEmbed(**patch_cfg)
+    wrapped_model.eval()
+    with RewriterContext({}, backend='ncnn'), torch.no_grad():
+        _, shape = wrapped_model(input)
+        assert shape[0] == patch_cfg['input_size'] / patch_cfg['stride']

tools/deploy.py

@@ -220,7 +220,9 @@ def main():
         from mmdeploy.apis.ncnn import is_available as is_available_ncnn
 
         if not is_available_ncnn():
-            logger.error('ncnn support is not available.')
+            logger.error('ncnn support is not available, please make sure \
+                1) `onnx2ncnn` existed in `PATH` 2) python import ncnn success'
+                         )
             exit(1)
 
         import mmdeploy.apis.ncnn as ncnn_api