feat(mgeconvert): support relu inplace and remove flatten before linear for caffe

Author: dingshaohua960303

bin/convert

@@ -123,6 +123,18 @@ def init(subparsers):
                 "tracedmodule_to_caffe": mgeconvert.tracedmodule_to_caffe,
                 "mge_to_caffe": mgeconvert.mge_to_caffe,
             }
+            convert_backend_map = {
+                "caffe": 1,
+                "1":1,
+                "snpe": 2,
+                "2":2,
+                "trt": 3,
+                "3":3,
+                "nnie": 4,
+                "4":4
+            }
+            assert args.convert_backend in convert_backend_map, "unsupported backend"
+            convert_backend = convert_backend_map[args.convert_backend]
             if "tracedmodule" in target:
                 converter_map[target](
                     args.input,
@@ -137,15 +149,15 @@ def init(subparsers):
                     split_conv_relu=args.split_conv_relu,
                     fuse_bn=args.fuse_bn,
                     quantize_file_path=args.quantize_file_path,
-                    convert_backend=args.convert_backend,
+                    convert_backend=convert_backend,
                 )
             else:
                 converter_map[target](
                     args.input,
                     prototxt=args.prototxt,
                     caffemodel=args.caffemodel,
                     outspec=outspec,
-                    convert_backend=args.convert_backend,
+                    convert_backend=convert_backend,
                 )
 
         def caffe_parser(subparsers):
@@ -240,9 +252,9 @@ def init(subparsers):
 
             p.add_argument(
                 "--convert_backend",
-                default=1,
-                type=int,
-                help="backend  1 for caffe, 2 for snpe, 3 for trt",
+                default='1',
+                type=str,
+                help="backend  '1' or 'caffe' for caffe, '2' or 'snpe' for snpe, '3' or 'trt' for trt, '4' or 'nnie' for nnie",
             )
 
         caffe_parser(subparsers)

mgeconvert/backend/ir_to_caffe/caffe_op.py

@@ -78,6 +78,7 @@ class BackEnd(IntEnum):
     CAFFE = 1
     SNPE = 2
     TRT = 3
+    NNIE = 4
 
 
 def isconst(x):
@@ -118,7 +119,10 @@ def _gen_layer(opr, etype, context, single_input=True, **kwargs):
         if single_input
         else list(map(context.get_blob_name, opr.inp_tensors))
     )
-    top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)]
+    if etype == "ReLU" and context.convert_backend == BackEnd.NNIE:
+        top = [context.set_blob_name(opr.out_tensors[0], opr.inp_tensors[0].name)]
+    else:
+        top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)]
     return cp.LayerParameter(
         bottom=bottom, top=top, name=opr.out_tensors[0].name, type=etype, **kwargs
     )
@@ -949,11 +953,8 @@ def _reshape(opr, context):
         tmp_shape = out_shape + (1,) * d
         bottom = [context.get_blob_name(opr.inp_tensors[0])]
         top = [context.set_blob_name(opr.out_tensors[0], opr.out_tensors[0].name)]
-        if context.convert_backend == BackEnd.CAFFE:
+        if context.convert_backend == BackEnd.NNIE:
             if inp_shape != tmp_shape:
-                logger.warning(
-                    "trt don't support flatten after reshape, but caffe support! please use BackEnd.TRT in tracedmodule_to_caffe by pass convert_backend:BackEnd=BackEnd.TRT"
-                )
                 param = cp.ReshapeParameter(shape=cp.BlobShape(dim=tmp_shape))
                 tmp = [bottom[0] + "tmp"]
                 context.add_layer(
@@ -976,18 +977,17 @@ def _reshape(opr, context):
                     flatten_param=param,
                 )
             )
-        elif context.convert_backend == BackEnd.TRT:
-            if inp_shape != tmp_shape:
-                param = cp.ReshapeParameter(shape=cp.BlobShape(dim=tmp_shape))
-                context.add_layer(
-                    cp.LayerParameter(
-                        bottom=bottom,
-                        top=top,
-                        name=opr.out_tensors[0].name,
-                        type="Reshape",
-                        reshape_param=param,
-                    )
+        else:
+            param = cp.ReshapeParameter(shape=cp.BlobShape(dim=out_shape))
+            context.add_layer(
+                cp.LayerParameter(
+                    bottom=bottom,
+                    top=top,
+                    name=opr.out_tensors[0].name,
+                    type="Reshape",
+                    reshape_param=param,
                 )
+            )
     else:
         logger.warning(
             "NNIE doesn't support this reshape Opr %s, inp_shape %s, out_shape %s, NNIE reshape only support C/H/W, not N!",

mgeconvert/converter_ir/ir_transform.py

@@ -32,6 +32,7 @@
     IdentityOpr,
     LeakyReluOpr,
     LinearOpr,
+    MatMulOpr,
     MulOpr,
     OpBase,
     PadOpr,
@@ -83,7 +84,8 @@ class TransformerRule(Enum):
     # for TFLite Converter
     SLICE_PARAMS_AS_INPUTS_AND_MAKE_SQUEEZE = 119
     RESIZE_PARAMS_AS_INPUT = 120
-    REPLACE_FLATTEN_TO_RESHAPE = 120.1
+    REMOVE_FLATTEN_BEFORE_LINEAR = 120.1
+    REPLACE_FLATTEN_TO_RESHAPE = 120.2
 
     # remove reshape
     REMOVE_RESHAPE_REALTED_OP = 121
@@ -1273,3 +1275,25 @@ def _expand_conv_relu(net: IRGraph):
             t.owner_opr = relu_op
 
         net.replace_op(opr, relu_op)
+
+
+@_register_tranformation_rule(TransformerRule.REMOVE_FLATTEN_BEFORE_LINEAR)
+def _remove_flatten(net: IRGraph):
+    for opr in net.all_oprs:
+        if not isinstance(opr, MatMulOpr):
+            continue
+        inp_oprs = net.find_inp_oprs(opr)
+
+        if (
+            isinstance(inp_oprs[0], OpBase)
+            and isinstance(inp_oprs[0], FlattenOpr)
+            and len(net.find_out_oprs(inp_oprs[0])) == 1
+            and net.find_out_oprs(inp_oprs[0])[0] == opr
+        ):
+            assert opr.inp_tensors[0].np_data is None
+            flatten_opr = inp_oprs[0]
+            opr.inp_tensors[0] = flatten_opr.inp_tensors[0]
+            flatten_opr.inp_tensors[0].user_opr.remove(flatten_opr)
+            flatten_opr.inp_tensors[0].user_opr.append(opr)
+            idx = net.all_oprs.index(flatten_opr)
+            net.delete_ops(idx)

mgeconvert/converters/tm_to_caffe.py

@@ -75,6 +75,11 @@ def tracedmodule_to_caffe(
             TransformerRule.FUSE_CONV_BN,
         ]
 
+    if convert_backend == BackEnd.NNIE:
+        transformer_options.extend(
+            [TransformerRule.REMOVE_FLATTEN_BEFORE_LINEAR,]
+        )
+
     if split_conv_relu:
         transformer_options += [TransformerRule.REMOVE_RELU]
     transformer = IRTransform(transformer_options)

test/traced_module/test_caffe.py

@@ -58,6 +58,7 @@ def _test_convert_result(
     split_conv_relu=False,
     fuse_bn=False,
     input_name="x",
+    convert_backend=1,
 ):
     tracedmodule_to_caffe(
         trace_module,
@@ -70,6 +71,7 @@ def _test_convert_result(
         param_fake_quant=param_fake_quant,
         split_conv_relu=split_conv_relu,
         fuse_bn=fuse_bn,
+        convert_backend=convert_backend,
     )
 
     caffe_net = caffe.Net(tmp_file + ".txt", tmp_file + ".caffemodel", caffe.TEST)
@@ -134,6 +136,12 @@ def test_linear():
     _test_convert_result(net.data, tm_module, mge_result, max_error)
 
 
+def test_flatten_linear():
+    net = LinearOpr("flatten")
+    tm_module, mge_result = get_traced_module(net, mge.tensor(net.data1))
+    _test_convert_result(net.data1, tm_module, mge_result, max_error, convert_backend=4)
+
+
 def test_linear_bn():
     net = LinearBnOpr()
     for _ in range(10):
@@ -249,6 +257,13 @@ def test_active(mode):
     _test_convert_result(net.data, tm_module, mge_result, max_error)
 
 
+@pytest.mark.parametrize("mode", ["relu",])
+def test_active_inplace(mode):
+    net = ActiveOpr(mode)
+    tm_module, mge_result = get_traced_module(net, mge.tensor(net.data))
+    _test_convert_result(net.data, tm_module, mge_result, max_error, convert_backend=4)
+
+
 @pytest.mark.parametrize("mode", ["max", "sum", "mean"])
 def test_reduce(mode):
     net = ReduceOpr(mode)

test/utils.py

@@ -120,16 +120,21 @@ def forward(self, x):
 
 
 class LinearOpr(M.Module):
-    def __init__(self):
+    def __init__(self, mode="normal"):
         super().__init__()
         self.data = np.random.random((10, 100)).astype(np.float32)
+        self.data1 = np.random.random((10, 10, 10)).astype(np.float32)
         self.linear = M.Linear(100, 200, bias=False)
         self.linear_bias = M.Linear(200, 200, bias=True)
         self.linear_bias.bias = mge.Parameter(
             np.random.random(self.linear_bias.bias.shape).astype(np.float32)
         )
+        self.mode = mode
 
     def forward(self, x):
+        if self.mode == "flatten":
+            x = F.flatten(x, 1)
+
         x = self.linear(x)
         x = self.linear_bias(x)
         return x