Merge pull request #443 from onnx/gs/code-reorg

reorg handlers in tfonnx.py into category based files

Author: guschmue

tests/run_pretrained_models.yaml

@@ -185,17 +185,6 @@ mobilenet_v1_75_192:
   outputs:
     - MobilenetV1/Predictions/Softmax:0
 
-tiny-yolo:
-  # works but local file
-  disabled: true
-  model: c:/src/darkflow/built_graph/tiny-yolo.pb
-  input_get: get_beach
-  inputs:
-    "input:0": [1, 416, 416, 3]
-  outputs:
-    - output:0
-  rtol: 0.6
-
 nasnet-a_mobile_224:
   # has only checkpoint format
   disabled: true

tests/test_graph.py

@@ -22,10 +22,12 @@
 from tf2onnx.graph import GraphUtil
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
 from tf2onnx.tfonnx import process_tf_graph
+from tf2onnx.handler import tf_op
+
 from common import get_test_config, unittest_main
 
 
-# pylint: disable=missing-docstring
+# pylint: disable=missing-docstring,unused-argument,unused-variable
 
 def onnx_to_graphviz(g, include_attrs=False):
     """Return dot for graph."""
@@ -331,10 +333,10 @@ def rewrite_test(g, ops):
                 'output [op_type=Identity] input1:0 -> Add input1:0 -> Add Add:0 -> output }',
                 onnx_to_graphviz(g))
 
-    def test_custom_op(self):
-        """Custom op test."""
+    def test_custom_op_depreciated(self):
+        """Custom op test using old depreciated api."""
 
-        def print_handler(ctx, node, name, args):  # pylint: disable=unused-argument
+        def print_handler(ctx, node, name, args):
             # replace tf.Print() with Identity
             #   T output = Print(T input, data, @list(type) U, @string message, @int first_n, @int summarize)
             # becomes:
@@ -360,6 +362,32 @@ def print_handler(ctx, node, name, args):  # pylint: disable=unused-argument
             self.assertEqual(g.opset, self.config.opset)
             self.assertEqual(g.extra_opset, [constants.TENSORFLOW_OPSET])
 
+    def test_custom_op(self):
+        """Custom op test."""
+
+        @tf_op("Print", onnx_op="Identity")
+        class Print:
+            @classmethod
+            def version_1(cls, ctx, node, **kwargs):
+                self.assertEqual(node.type, "Identity")
+                node.domain = constants.TENSORFLOW_OPSET.domain
+                del node.input[1:]
+                return node
+
+        with tf.Session() as sess:
+            x = tf.placeholder(tf.float32, [2, 3], name="input1")
+            x_ = tf.Print(x, [x], "hello")
+            _ = tf.identity(x_, name="output")
+            g = process_tf_graph(sess.graph,
+                                 opset=self.config.opset,
+                                 extra_opset=[constants.TENSORFLOW_OPSET])
+            self.assertEqual(
+                'digraph { input1 [op_type=Placeholder shape="[2, 3]"] Print [domain="ai.onnx.converters.tensorflow" '
+                'op_type=Identity] output [op_type=Identity] input1:0 -> Print Print:0 -> output }',
+                onnx_to_graphviz(g))
+            self.assertEqual(g.opset, self.config.opset)
+            self.assertEqual(g.extra_opset, [constants.TENSORFLOW_OPSET])
+
     def test_extra_opset(self):
         extra_opset = [
             utils.make_opsetid(constants.MICROSOFT_DOMAIN, 1),

tf2onnx/constants.py

@@ -27,3 +27,8 @@
 TARGET_CAFFE2 = "caffe2"
 POSSIBLE_TARGETS = [TARGET_RS4, TARGET_RS5, TARGET_RS6, TARGET_CAFFE2]
 DEFAULT_TARGET = []
+
+NCHW_TO_NHWC = [0, 2, 3, 1]
+NHWC_TO_NCHW = [0, 3, 1, 2]
+HWCN_TO_NCHW = [3, 2, 0, 1]
+NCHW_TO_HWCN = [2, 3, 1, 0]

tf2onnx/custom_opsets/__init__.py

@@ -3,8 +3,3 @@
 """ custom tf2onnx mapping functions. """
 
 from . import ms
-from .. import constants
-
-DOMAIN_OPSETS = {
-    constants.MICROSOFT_DOMAIN: ms.OPSETS
-}

tf2onnx/custom_opsets/ms.py

@@ -4,24 +4,15 @@
 
 from onnx.onnx_pb import TensorProto
 from tf2onnx import constants, utils
-from tf2onnx.function.range import make_range_const
+from tf2onnx.handler import tf_op
+from tf2onnx.onnx_opset import controlflow
 
 
-# pylint: disable=unused-argument
-
-def range_op1(ctx, node, name, args):
-    """Range."""
-    # T range = Range(T start, T limit, T delta)
-    dtype = node.get_attr_int("Tidx")
-    shape = node.output_shapes[0]
-    utils.make_sure(dtype is not None, "Tidx of %s is None", node.name)
-    ctx.remove_node(node.name)
-    make_range(ctx, node.input[0], node.input[1], node.input[2], node.output[0], name, shape, dtype)
-
+# pylint: disable=unused-argument,missing-docstring
 
 def make_range(ctx, start, limit, delta, output, scope_name, shape, dtype):
     if all(ctx.get_node_by_output(n).is_const() for n in [start, limit, delta]) is True:
-        make_range_const(ctx, start, limit, delta, output, scope_name, shape, dtype)
+        controlflow.make_range_const(ctx, start, limit, delta, output, scope_name, shape, dtype)
     else:
         _make_range_non_const(ctx, start, limit, delta, output, scope_name, shape, dtype)
 
@@ -34,10 +25,14 @@ def _make_range_non_const(ctx, start, limit, delta, output, scope_name, shape, d
                   domain=constants.MICROSOFT_DOMAIN)
 
 
-_OPSET_1 = {
-    "Range": (range_op1, []),
-}
-
-OPSETS = [
-    (1, _OPSET_1),
-]
+@tf_op("Range", domain=constants.MICROSOFT_DOMAIN)
+class Range:
+    @classmethod
+    def version_1(cls, ctx, node, **kwargs):
+        """Range."""
+        # T range = Range(T start, T limit, T delta)
+        dtype = node.get_attr_int("Tidx")
+        shape = node.output_shapes[0]
+        utils.make_sure(dtype is not None, "Tidx of %s is None", node.name)
+        ctx.remove_node(node.name)
+        make_range(ctx, node.input[0], node.input[1], node.input[2], node.output[0], node.name, shape, dtype)