improvement(ViT): use Crop to subtitude Gather (open-mmlab#477)

* improvement(ViT): use Crop to subtitude Gather

* fix(CI): code format

* fix(pytorch/ops/linear.py): bias maybe None

* fix(test/test_pytorch_ops.py): op_type error

* fix(test): pytest error

* fix(test): torch version 1.8

csrc/backend_ops/ncnn/onnx2ncnn/fuse_pass.cpp

@@ -1,6 +1,44 @@
 // Copyright (c) OpenMMLab. All rights reserved.
 #include "fuse_pass.h"
 
+void fuse_rewrite_gather(onnx::GraphProto* mutable_graph,
+                         std::map<std::string, onnx::TensorProto>& weights,
+                         std::map<std::string, int>& node_reference,
+                         std::set<std::string>& blob_names, int& reduced_node_count) {
+  const int node_count = mutable_graph->node_size();
+  for (int i = 0; i < node_count; ++i) {
+    onnx::NodeProto* gather = mutable_graph->mutable_node(i);
+    if (gather->op_type() != "Gather") {
+      continue;
+    }
+    auto indices = get_node_attr_from_input_ai(weights[gather->input(1)]);
+    if (indices.size() != 1) {
+      continue;
+    }
+
+    {
+      // reconstruct node connections
+      node_reference[gather->input(1)] -= 1;
+      std::string origin_inp = gather->input(0);
+      gather->clear_input();
+      gather->add_input(origin_inp);
+    }
+
+    {
+      // update axis, starts and ends
+      int axis = get_node_attr_i(*gather, "axis", 1) - 1;
+
+      gather->set_op_type("Crop");
+      gather->clear_attribute();
+
+      int indice = indices[0];
+      set_node_attr_ai(*gather, "starts", std::vector<int>{indice});
+      set_node_attr_ai(*gather, "ends", std::vector<int>{indice + 1});
+      set_node_attr_ai(*gather, "axis", std::vector<int>{axis});
+    }
+  }
+}
+
 void fuse_weight_reshape(onnx::GraphProto* mutable_graph,
                          std::map<std::string, onnx::TensorProto>& weights,
                          std::map<std::string, int>& node_reference,

csrc/backend_ops/ncnn/onnx2ncnn/fuse_pass.h

@@ -4,6 +4,11 @@
 #include "shape_inference.h"
 #include "utils.h"
 
+void fuse_rewrite_gather(onnx::GraphProto* mutable_graph,
+                         std::map<std::string, onnx::TensorProto>& weights,
+                         std::map<std::string, int>& node_reference,
+                         std::set<std::string>& blob_names, int& reduced_node_count);
+
 void fuse_weight_reshape(onnx::GraphProto* mutable_graph,
                          std::map<std::string, onnx::TensorProto>& weights,
                          std::map<std::string, int>& node_reference,

csrc/backend_ops/ncnn/onnx2ncnn/onnx2ncnn.cpp

@@ -229,6 +229,7 @@ int main(int argc, char** argv) {
     fuse_multiheadattention(mutable_graph, weights, node_reference, blob_names, reduced_node_count);
     fuse_binaryop_with_scalar(mutable_graph, weights, node_reference, blob_names,
                               reduced_node_count);
+    fuse_rewrite_gather(mutable_graph, weights, node_reference, blob_names, reduced_node_count);
   }
 
   // reduce common const weight node_reference
@@ -623,6 +624,8 @@ int main(int argc, char** argv) {
       }
     } else if (op == "Cos") {
       fprintf(pp, "%-16s", "UnaryOp");
+    } else if (op == "Crop") {
+      fprintf(pp, "%-16s", "Crop");
     } else if (op == "DepthToSpace") {
       fprintf(pp, "%-16s", "PixelShuffle");
     } else if (op == "DetectionOutput") {
@@ -1196,6 +1199,22 @@ int main(int argc, char** argv) {
     } else if (op == "Cos") {
       int op_type = 10;
       fprintf(pp, " 0=%d", op_type);
+    } else if (op == "Crop") {
+      auto starts = get_node_attr_ai(node, "starts");
+      fprintf(pp, " -23309=%zu", starts.size());
+      for (size_t j = 0; j < starts.size(); ++j) {
+        fprintf(pp, ",%i", starts[j]);
+      }
+      auto ends = get_node_attr_ai(node, "ends");
+      fprintf(pp, " -23310=%zu", ends.size());
+      for (size_t j = 0; j < ends.size(); ++j) {
+        fprintf(pp, ",%i", ends[j]);
+      }
+      auto axis = get_node_attr_ai(node, "axis");
+      fprintf(pp, " -23311=%zu", axis.size());
+      for (size_t j = 0; j < axis.size(); ++j) {
+        fprintf(pp, ",%i", axis[j]);
+      }
     } else if (op == "DepthToSpace") {
       // pixelshuffle
       int scale_factor = get_node_attr_i(node, "blocksize", 1);
@@ -1287,7 +1306,7 @@ int main(int argc, char** argv) {
       }
       fprintf(pp, " 0=%d", axis);
     } else if (op == "Gelu") {
-      fprintf(pp, " 0=0");
+      fprintf(pp, " 0=1");
     } else if (op == "Gemm") {
       float alpha = get_node_attr_f(node, "alpha", 1.f);
       float beta = get_node_attr_f(node, "beta", 1.f);

csrc/backend_ops/ncnn/onnx2ncnn/shape_inference.cpp

@@ -2,6 +2,8 @@
 
 #include "shape_inference.h"
 
+#include <algorithm>
+
 /**
  * @brief query output shape of target node
  *

csrc/backend_ops/ncnn/onnx2ncnn/shape_inference.h

@@ -1,7 +1,6 @@
 // Copyright (c) OpenMMLab. All rights reserved.
 
 #pragma once
-#include <algorithm>
 
 #include "utils.h"
 

mmdeploy/pytorch/functions/linear.py

@@ -25,7 +25,7 @@ def linear__ncnn(
 
     dim = input.dim()
 
-    if dim == 2:
+    if dim == 2 or dim == 3 and input.shape[0] == 1:
         return origin_func(input, weight, bias)
     else:
         out = origin_func(input, weight)

mmdeploy/pytorch/ops/__init__.py

@@ -8,6 +8,7 @@
 from .hardsigmoid import hardsigmoid__default
 from .instance_norm import instance_norm__tensorrt
 from .layer_norm import layer_norm__ncnn
+from .linear import linear__ncnn
 from .lstm import generic_rnn__ncnn
 from .squeeze import squeeze__default
 
@@ -16,5 +17,5 @@
     'adaptive_avg_pool3d__default', 'grid_sampler__default',
     'hardsigmoid__default', 'instance_norm__tensorrt', 'generic_rnn__ncnn',
     'squeeze__default', 'adaptive_avg_pool2d__ncnn', 'gelu__ncnn',
-    'layer_norm__ncnn'
+    'layer_norm__ncnn', 'linear__ncnn'
 ]

mmdeploy/pytorch/ops/linear.py

@@ -0,0 +1,44 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+# Modified from:
+# https://github.com/pytorch/pytorch/blob/9ade03959392e5a90b74261012de1d806cab2253/torch/onnx/symbolic_opset9.py
+from torch.onnx.symbolic_helper import parse_args
+
+from mmdeploy.core import SYMBOLIC_REWRITER
+from mmdeploy.utils import Backend
+
+
+@parse_args('v', 'v', 'f', 'f', 'i', 'i')
+def linear_no_bias(g, input, weight):
+    """Symbolic function for `linear` without bias.
+
+    PyTorch `nn.Linear` will be exported as ONNX node 'Gemm'.
+    """
+    return g.op(
+        'Gemm', input, weight, alpha_f=1.0, beta_f=1.0, transA_i=0, transB_i=1)
+
+
+@parse_args('v', 'v', 'v', 'f', 'f', 'i', 'i')
+def linear_normal(g, input, weight, bias):
+    """Symbolic function for `linear`.
+
+    PyTorch `nn.Linear` will be exported as ONNX node 'Gemm'.
+    """
+    return g.op(
+        'Gemm',
+        input,
+        weight,
+        bias,
+        alpha_f=1.0,
+        beta_f=1.0,
+        transA_i=0,
+        transB_i=1)
+
+
+@SYMBOLIC_REWRITER.register_symbolic(
+    'linear', is_pytorch=True, backend=Backend.NCNN.value)
+def linear__ncnn(ctx, g, input, weight, bias):
+    """Support export linear This rewrite enable export Gemm."""
+    if bias is None:
+        return linear_no_bias(g, input, weight)
+    else:
+        return linear_normal(g, input, weight, bias)

tests/test_pytorch/test_pytorch_ops.py

@@ -127,6 +127,41 @@ def test_instance_norm():
     assert nodes[4].domain == 'mmdeploy'
 
 
+@pytest.mark.usefixtures('prepare_symbolics_ncnn')
+class TestLinear:
+
+    def check(self, nodes):
+        print(nodes)
+
+        from packaging.version import parse as version_parse
+        version = version_parse(torch.__version__)
+        target = 'Gemm'
+        if version.major <= 1 and version.minor <= 8:
+            target = 'MatMul'
+        exist = False
+        for node in nodes:
+            if node.op_type == target:
+                exist = True
+                break
+
+        assert exist is True
+
+    def test_normal(self):
+        x = torch.rand(1, 2, 3)
+        w = torch.rand(2, 3)
+        bias = torch.rand(2)
+        model = OpModel(torch.nn.functional.linear, w, bias).eval()
+        nodes = get_model_onnx_nodes(model, x)
+        self.check(nodes)
+
+    def test_no_bias(self):
+        x = torch.rand(1, 2, 3)
+        w = torch.rand(2, 3)
+        model = OpModel(torch.nn.functional.linear, w).eval()
+        nodes = get_model_onnx_nodes(model, x)
+        self.check(nodes)
+
+
 @pytest.mark.usefixtures('prepare_symbolics')
 class TestSqueeze: