add groupconv2d

hanjr92 · hanjr92 · commit 2b08d2454504 · 2022-07-22T10:24:31.000+08:00
diff --git a/tests/test_groupconv2d.py b/tests/test_groupconv2d.py
@@ -0,0 +1,49 @@
+#! /usr/bin/python
+# -*- coding: utf-8 -*-
+
+import os
+# os.environ["TL_BACKEND"] = 'tensorflow'
+# os.environ["TL_BACKEND"] = 'paddle'
+os.environ["TL_BACKEND"] = 'torch'
+# os.environ["TL_BACKEND"] = 'mindspore'
+import tensorlayerx as tlx
+from tensorlayerx.nn import Module
+from tensorlayerx.nn import GroupConv2d
+from tlx2onnx.main import export
+import onnxruntime as rt
+import numpy as np
+
+
+############################################ test 2d ###########################################################
+class CNN(Module):
+
+    def __init__(self):
+        super(CNN, self).__init__()
+        # weights init
+        W_init = tlx.nn.initializers.truncated_normal(stddev=5e-2)
+        b_init2 = tlx.nn.initializers.constant(value=0.1)
+        self.conv1 = GroupConv2d(
+            36, (5, 5), (1, 1), n_group=3, padding=(2,2), W_init=W_init, b_init=b_init2, name='conv1',
+            in_channels=3, data_format='channels_last', act = tlx.nn.ReLU
+        )
+    def forward(self, x):
+        z = self.conv1(x)
+        return z
+
+net = CNN()
+input = tlx.nn.Input(shape=(1, 10, 10, 3))
+net.set_eval()
+output = net(input)
+print("groupconv2d tlx output", output)
+onnx_model = export(net, input_spec=input, path='groupconv2d_model.onnx')
+
+# Infer Model
+sess = rt.InferenceSession('groupconv2d_model.onnx')
+
+input_name = sess.get_inputs()[0].name
+output_name = sess.get_outputs()[0].name
+
+input_data = np.array(input, dtype=np.float32)
+
+result = sess.run([output_name], {input_name: input_data})
+print("groupconv2d onnx output", result)
diff --git a/tlx2onnx/op_mapper/nn/__init__.py b/tlx2onnx/op_mapper/nn/__init__.py
@@ -23,4 +23,5 @@
 from .stack import *
 from .subpixelconv import *
 from .mask_conv import *
+from .group_conv import *
 
diff --git a/tlx2onnx/op_mapper/nn/group_conv.py b/tlx2onnx/op_mapper/nn/group_conv.py
@@ -0,0 +1,255 @@
+#! /usr/bin/python
+# -*- coding: utf-8 -*-
+import numpy as np
+from onnx import helper, numpy_helper
+from collections import OrderedDict
+import tensorlayerx as tlx
+from tlx2onnx.op_mapper.datatype_mapping import NP_TYPE_TO_TENSOR_TYPE
+from tlx2onnx.op_mapper.op_mapper import OpMapper
+from tlx2onnx.common import make_node
+from tlx2onnx.common import make_shape_channels_first, get_channels_first_permutation,tlx_act_2_onnx,get_channels_last_permutation
+
+
+def convert_padding(padding, input_shape, output_shape, kernel_shape, strides, dilations, spatial, data_format):
+    if isinstance(padding, str):
+        if padding == "SAME":
+            pads = [0] * (spatial * 2)
+            if data_format == "channels_last":
+                input_shape = make_shape_channels_first(input_shape)
+                output_shape = make_shape_channels_first(output_shape)
+
+            if any(input_shape[i + 2] == -1 or output_shape[i + 2] == -1 for i in range(spatial)):
+
+                auto_pad = "SAME_UPPER"
+
+                return  auto_pad
+
+            for i in range(spatial):
+                pad = (
+                    (output_shape[i + 2] - 1) * strides[i]
+                    + dilations[i] * (kernel_shape[i] - 1) + 1
+                    - input_shape[i + 2]
+                )
+                pad = max(pad, 0)
+                pads[i] = pad // 2
+                pads[i + spatial] = pad - pad // 2
+
+            return pads
+
+        elif padding == "VALID":
+            auto_pad = "VALID"
+            return auto_pad
+    elif isinstance(padding, int):
+        pads = [padding] * spatial * 2
+        return pads
+    elif isinstance(padding, tuple):
+        return list(padding) * 2
+
+def convert_w(w, data_format, spatial, w_name):
+    w = tlx.convert_to_numpy(w)
+    if tlx.BACKEND == 'tensorflow':
+        if spatial == 2:
+            w = np.transpose(w, axes=[3, 2, 0, 1])
+        elif spatial == 1:
+            w = np.transpose(w, axes=[2, 1, 0])
+        elif spatial == 3:
+            w = np.transpose(w, axes=[4, 3, 0, 1, 2])
+        return numpy_helper.from_array(w, name=w_name)
+    elif tlx.BACKEND == 'mindspore':
+        if spatial == 2 and data_format == 'channels_last':
+            w = np.transpose(w, axes=[3, 0, 1, 2])
+            return numpy_helper.from_array(w, name=w_name)
+    return numpy_helper.from_array(w, name=w_name)
+
+def convert_b(b, b_name):
+    b = tlx.convert_to_numpy(b)
+    return numpy_helper.from_array(b, name=b_name)
+
+
+@OpMapper(["GroupConv2d"])
+class Conv():
+    # suppport v1-v13
+
+    @classmethod
+    def any_version(cls, node, opset, **kwargs):
+        """
+        Parameters
+        ----------
+        node:node dict {node: node,
+                            in_tensors: node inputs,
+                            out_tensors: node outputs,
+                            in_nodes_name: node inputs name,
+                            out_nodes_name: node outputs name}
+        Returns
+        -------
+        """
+        Op_name = 'Conv'
+        onnx_node, onnx_value, onnx_init = [], [], []
+        attr_dict = OrderedDict()
+
+        #### get data_type
+        data_type = node['dtype']
+        tensor_type = NP_TYPE_TO_TENSOR_TYPE[data_type]
+        #### get in_node_name out_node_nmae
+        x_name = node['in_nodes_name'][0]
+        out_name = node['out_nodes_name'][0]
+        x_shape = node['in_tensors'][0]
+        out_shape = node['out_tensors'][0]
+
+        #### get cur_node_layer node_index
+        layer = node['node'].layer
+        layer_type = layer.__class__.__name__
+        spatial = int(layer_type[-2])
+        node_name = layer.name
+        #### get layer_param
+        layer_param = layer.all_weights
+
+        #### get layer_act_type
+        layer_act = layer.act.__class__.__name__
+
+        #### conv inputs
+        w = None
+        b = None
+        if len(layer_param) == 1:
+            w = layer_param[0]
+        elif len(layer_param) == 2:
+            w = layer_param[0]
+            b = layer_param[1]
+
+        #### insert conv attr
+        kernel_size = node['attr']['kernel_size']
+        if isinstance(kernel_size, int):
+            kernel_size = [kernel_size]
+        attr_dict["kernel_shape"] = kernel_size
+        dilations = node['attr']['dilation']
+        if isinstance(dilations, int):
+            dilations = [dilations,]
+        attr_dict["dilations"] = dilations
+        strides = node['attr']['stride']
+        if isinstance(strides, int):
+            strides = [strides]
+        attr_dict["strides"] = strides
+        data_format = node['attr']['data_format']
+        paddding = node['attr']['padding']
+        attr_dict["group"] = layer.n_group
+        attr_dict["outputs"] = [out_name]
+
+        ####convert padding
+        pads = convert_padding(
+            paddding, x_shape, out_shape, attr_dict["kernel_shape"], attr_dict["strides"],
+            attr_dict["dilations"], spatial, data_format
+        )
+        if isinstance(pads, str):
+            attr_dict["auto_pad"] = pads
+        else:
+            attr_dict["pads"] = pads
+
+        if data_format == 'channels_last':
+            permutation = get_channels_first_permutation(spatial)
+            x_shape_t = make_shape_channels_first(x_shape)
+            # insert transpose op: NHWC -> NCHW
+            transpose_value = helper.make_tensor_value_info(x_name+'_t', tensor_type, shape=x_shape_t)
+            onnx_value.append(transpose_value)
+            transpose_node, out = make_node('Transpose', inputs=[x_name], outputs=[x_name+'_t'], perm = permutation)
+            onnx_node.append(transpose_node)
+            # convert w
+            w_name = node_name + '_w'
+            w_init = convert_w(w, data_format, spatial, w_name)
+            onnx_init.append(w_init)
+            attr_dict["inputs"] = [out, w_name]
+
+            #### convert b
+            if b is not None:
+                b_name = node_name + '_b'
+                b_init = convert_b(b, b_name)
+                onnx_init.append(b_init)
+                attr_dict["inputs"] = [out, w_name, b_name]
+
+            attr_dict["outputs"] = [out + "_t"]
+            conv_node, out = make_node(Op_name, **attr_dict)
+            onnx_node.append(conv_node)
+            out_shape_t = make_shape_channels_first(out_shape)
+            conv_value = helper.make_tensor_value_info(out, tensor_type, shape=out_shape_t)
+            onnx_value.append(conv_value)
+            # insert transpose op: NCHW -> NHWC  and  insert act node
+
+            if layer_act != 'NoneType':
+                act_convert = tlx_act_2_onnx[layer_act]
+                act_input = out_name + "_act"
+                act_out = out_name
+                # insert transpose op
+                permutation = get_channels_last_permutation(spatial)
+                transpose_node, out = make_node('Transpose', inputs=[out], outputs=[act_input], perm = permutation)
+                onnx_node.append(transpose_node)
+                transpose_value = helper.make_tensor_value_info(act_input, tensor_type, shape = out_shape)
+                onnx_value.append(transpose_value)
+                # 如果layer存在act，需要新增一个act node 和 对应act输入的act input info， 并且要更新 conv的outputs 为 act的inputs， 此时act的outputs是整个layer的outputs
+                act_node, _ = act_convert([out], [act_out])
+                act_input_value_info = helper.make_tensor_value_info(act_out, tensor_type, out_shape)
+                onnx_value.append(act_input_value_info)
+                onnx_node.append(act_node)
+                return onnx_node, onnx_value, onnx_init
+            else:
+                permutation = get_channels_last_permutation(spatial)
+                transpose_node, out = make_node('Transpose', inputs=[out], outputs=[out_name], perm=permutation)
+                onnx_node.append(transpose_node)
+                transpose_value = helper.make_tensor_value_info(out_name, tensor_type, shape=out_shape)
+                onnx_value.append(transpose_value)
+                return onnx_node, onnx_value, onnx_init
+
+
+        elif data_format == 'channels_first':
+
+            #### convert w
+            w_name = node_name + '_w'
+            w_init = convert_w(w, data_format, spatial, w_name)
+            onnx_init.append(w_init)
+            attr_dict["inputs"] = [x_name, w_name]
+
+            #### convert b
+            if b is not None:
+                b_name = node_name + '_b'
+                b_init = convert_b(b, b_name)
+                onnx_init.append(b_init)
+                attr_dict["inputs"] = [x_name, w_name, b_name]
+
+            #### make act node
+            if layer_act != 'NoneType':
+                act_convert = tlx_act_2_onnx[layer_act]
+                act_input = out_name + "_act"
+                act_out = out_name
+                attr_dict["outputs"] = [act_input]
+                conv_node, out = make_node(Op_name, **attr_dict)
+                onnx_node.append(conv_node)
+                conv_value = helper.make_tensor_value_info(out, tensor_type, shape = out_shape)
+                onnx_value.append(conv_value)
+                #insert act node
+                act_node, out = act_convert([act_input], [act_out])
+                act_input_value_info = helper.make_tensor_value_info(out, tensor_type, out_shape)
+                onnx_value.append(act_input_value_info)
+                onnx_node.append(act_node)
+                return onnx_node, onnx_value, onnx_init
+            else:
+                conv_node, out = make_node(Op_name, **attr_dict)
+                onnx_node.append(conv_node)
+                conv_value = helper.make_tensor_value_info(out, tensor_type, out_shape)
+                onnx_value.append(conv_value)
+                return onnx_node, onnx_value, onnx_init
+        else:
+            raise ValueError("Only support 'channels_first' or 'channels_last' data_format mode, but got {}.".format(data_format))
+
+    @classmethod
+    def version_1(cls, node, **kwargs):
+
+        return cls.any_version(node, 1, **kwargs)
+
+
+    @classmethod
+    def version_11(cls, node, **kwargs):
+
+        return cls.any_version( node, 11, **kwargs)
+
+    @classmethod
+    def version_13(cls, node, **kwargs):
+
+        return cls.any_version(node, 13, **kwargs)
