add depthtospace converter

Author: hanjr92

tests/test_subpixelconv.py

@@ -0,0 +1,40 @@
+#! /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 SubpixelConv2d
+from tlx2onnx.main import export
+import onnxruntime as rt
+import numpy as np
+tlx.set_seed(42)
+
+class conv(Module):
+
+    def __init__(self):
+        super(conv, self).__init__()
+        self.conv = SubpixelConv2d(scale=2, data_format="channels_last")
+
+    def forward(self, x):
+
+        x = self.conv(x)
+        return x
+
+model = conv()
+input = tlx.nn.Input(shape= (1, 2, 2, 8), init=tlx.nn.initializers.HeUniform())
+model.set_eval()
+output = model(input)
+print("SubpixelConv2d tlx output", output)
+onnx_model = export(model, input_spec=input, path='SubpixelConv2d.onnx', opset_version = 11)
+
+sess = rt.InferenceSession('SubpixelConv2d.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("SubpixelConv2d onnx output", result)

tlx2onnx/op_mapper/nn/__init__.py

@@ -20,4 +20,6 @@
 from .noise import *
 from .padding import *
 from .scale import *
-from .stack import *
+from .stack import *
+from .subpixelconv import *
+

tlx2onnx/op_mapper/nn/subpixelconv.py

@@ -0,0 +1,120 @@
+#! /usr/bin/python
+# -*- coding: utf-8 -*-
+import numpy as np
+import onnx
+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,get_channels_last_permutation
+
+@OpMapper(["SubpixelConv2d"])
+class SubpixelConv():
+
+    @classmethod
+    def version_1(cls, node, **kwargs):
+        onnx_node = []
+        onnx_value = []
+        onnx_init = []
+
+        op_type = "DepthToSpace"
+        attr_dict = OrderedDict()
+        # 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 data_type
+        data_type = node['dtype']
+        tensor_type = NP_TYPE_TO_TENSOR_TYPE[data_type]
+
+        # get cur_node_layer node_index
+        layer = node['node'].layer
+        layer_name = layer.__class__.__name__
+        spatial = int(layer_name[-2])
+
+        # get layer attr
+        scale = layer.scale
+        data_format = layer.data_format
+        attr_dict["blocksize"] = scale
+
+        if data_format == "channels_last":
+            permutation = get_channels_first_permutation(spatial)
+            x_shape_t = make_shape_channels_first(x_shape)
+            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)
+            depth_to_space, out = make_node(op_type, inputs=[out], outputs=[out + '_t'], **attr_dict)
+            onnx_node.append(depth_to_space)
+            permutation = get_channels_last_permutation(spatial)
+            transpose_node, out = make_node('Transpose', inputs=[out], outputs=[out_name], perm=permutation)
+            onnx_node.append(transpose_node)
+            return onnx_node, onnx_value, onnx_init
+
+        elif data_format == 'channels_first':
+
+            depth_to_space, out = make_node(op_type, inputs=[x_name], outputs=[out_name], **attr_dict)
+            onnx_node.append(depth_to_space)
+            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_11(cls, node, **kwargs):
+        onnx_node = []
+        onnx_value = []
+        onnx_init = []
+
+        op_type = "DepthToSpace"
+        attr_dict = OrderedDict()
+        # 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 data_type
+        data_type = node['dtype']
+        tensor_type = NP_TYPE_TO_TENSOR_TYPE[data_type]
+
+        # get cur_node_layer node_index
+        layer = node['node'].layer
+        layer_name = layer.__class__.__name__
+        spatial = int(layer_name[-2])
+
+        # get layer attr
+        scale = layer.scale
+        data_format = layer.data_format
+        attr_dict["blocksize"] = scale
+        if tlx.BACKEND in ["tensorflow", "mindspore"]:
+            attr_dict["mode"] = "DCR"
+        elif tlx.BACKEND in ["torch", "paddle"]:
+            attr_dict["mode"] = "CRD"
+
+        if data_format == "channels_last":
+            permutation = get_channels_first_permutation(spatial)
+            x_shape_t = make_shape_channels_first(x_shape)
+            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)
+            depth_to_space, out = make_node(op_type, inputs=[out], outputs=[out + '_t'], **attr_dict)
+            onnx_node.append(depth_to_space)
+            permutation = get_channels_last_permutation(spatial)
+            transpose_node, out = make_node('Transpose', inputs=[out], outputs=[out_name], perm=permutation)
+            onnx_node.append(transpose_node)
+            return onnx_node, onnx_value, onnx_init
+
+        elif data_format == 'channels_first':
+
+            depth_to_space, out = make_node(op_type, inputs=[x_name], outputs=[out_name], **attr_dict)
+            onnx_node.append(depth_to_space)
+            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))