Merge pull request #70 from onnx/gs/py27

changes to make it work on py2.7

Author: guschmue

setup.py

@@ -74,7 +74,7 @@ def run(self):
     version=VersionInfo.version,
     description='Tensorflow to ONNX converter',
     setup_requires=['pytest-runner'],
-    tests_require=['pytest', 'pytest-cov', 'psutil', 'graphviz', 'pyyaml'],
+    tests_require=['requests', 'pytest', 'pytest-cov', 'psutil', 'graphviz', 'pyyaml'],
     cmdclass=cmdclass,
     packages=find_packages(),
     author='[email protected]',

tests/run_pretrained_models.py

@@ -1,13 +1,16 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT license.
 
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
 import argparse
 import os
 import tarfile
 import time
 import tempfile
-import urllib
-import urllib.request
+import requests
 import zipfile
 
 import PIL.Image
@@ -136,7 +139,11 @@ def download_file(self):
         os.makedirs(dir_name, exist_ok=True)
         fpath = os.path.join(dir_name, fname)
         if not os.path.exists(fpath):
-            urllib.request.urlretrieve(url, fpath)
+            response = requests.get(url)
+            if response.status_code not in [200]:
+                response.raise_for_status()
+            with open(fpath, "wb") as f:
+                f.write(response.content)
         model_path = os.path.join(dir_name, self.local)
         if not os.path.exists(model_path):
             if ftype == 'tgz':

tests/test_backend.py

@@ -1,6 +1,9 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT license.
 
+from __future__ import division
+from __future__ import print_function
+
 import argparse
 import os
 import sys
@@ -325,7 +328,6 @@ def test_conv2d_6(self):
         expected, actual = self._conv_test(x_val, kernel_val, strides=strides, padding="VALID")
         self.assertAllClose(expected, actual, rtol=1e-05)
 
-
     def test_conv2d_7(self):
         x_shape = [1, 35, 35, 288]  # out: [1, 17, 17, 384]
         kernel_shape = [3, 3, 288, 384]
@@ -567,7 +569,7 @@ def test_logicaland(self):
         x2 = tf.placeholder(tf.bool, [2, 2], name=_TFINPUT1)
         mi = tf.logical_and(x1, x2)
         output = tf.identity(mi, name=_TFOUTPUT)
-        actual, expected = self._run(output, {x1: x_val1, x2: x_val2}, {_INPUT: x_val1, _INPUT1: x_val2,})
+        actual, expected = self._run(output, {x1: x_val1, x2: x_val2}, {_INPUT: x_val1, _INPUT1: x_val2})
         self.assertAllClose(expected, actual)
 
     def test_greater(self):
@@ -577,7 +579,7 @@ def test_greater(self):
         x2 = tf.placeholder(tf.float32, [2, 2], name=_TFINPUT1)
         mi = tf.greater(x1, x2)
         output = tf.identity(mi, name=_TFOUTPUT)
-        actual, expected = self._run(output, {x1: x_val1, x2: x_val2}, {_INPUT: x_val1, _INPUT1: x_val2,})
+        actual, expected = self._run(output, {x1: x_val1, x2: x_val2}, {_INPUT: x_val1, _INPUT1: x_val2})
         self.assertAllClose(expected, actual)
 
     def test_sequeeze(self):
@@ -780,9 +782,7 @@ def test_reducemean(self):
     @unittest.skip
     def test_slice1(self):
         # FIXME: only 1 dimension supported by caffe2 and msrt
-        x_val = np.array([[[1, 1, 1], [2, 2, 2]],
-                       [[3, 3, 3], [4, 4, 4]],
-                       [[5, 5, 5], [6, 6, 6]]], dtype=np.float32)
+        x_val = np.array([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]], [[5, 5, 5], [6, 6, 6]]], dtype=np.float32)
         t1 = tf.constant([1, 0, 0], dtype=tf.int32)
         t2 = tf.constant([1, 1, 3], dtype=tf.int32)
         x0 = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
@@ -920,7 +920,7 @@ def test_topk(self):
         x_val = np.arange(3*2*3).astype("float32")
         x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
         values, indices = tf.nn.top_k(x, 5, sorted=True)
-        output  = tf.identity(values, name=_TFOUTPUT)
+        output = tf.identity(values, name=_TFOUTPUT)
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
@@ -949,7 +949,7 @@ def test_space_to_depth(self):
         x_val = make_xval([1, 2, 2, 1])
         x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
         x_ = tf.space_to_depth(x, block_size=2)
-        output  = tf.identity(x_, name=_TFOUTPUT)
+        output = tf.identity(x_, name=_TFOUTPUT)
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
@@ -958,7 +958,7 @@ def test_addn(self):
         x_val = np.arange(3*2*3).astype("float32")
         x = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
         x_ = tf.add_n([x, x, x])
-        output  = tf.identity(x_, name=_TFOUTPUT)
+        output = tf.identity(x_, name=_TFOUTPUT)
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
         self.assertAllClose(expected, actual)
 
@@ -1020,7 +1020,7 @@ def test_batchnorm(self):
         x_shape = [1, 28, 28, 2]
         x_dtype = np.float32
         scale_dtype = np.float32
-        scale_shape =  [2]
+        scale_shape = [2]
         # only nhwc is support on cpu for tensorflow
         data_format = "NHWC"
         x_val = np.random.random_sample(x_shape).astype(x_dtype)

tests/test_graph.py

@@ -1,6 +1,9 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT license.
 
+from __future__ import division
+from __future__ import print_function
+
 import unittest
 from collections import namedtuple
 
@@ -116,7 +119,7 @@ def test_add(self):
             _ = tf.identity(x_, name="output")
             g = process_tf_graph(sess.graph)
             self.assertEqual(
-                'digraph { Add [op_type=Add] output [op_type=Identity] input1:0 -> Add input2:0 -> ' \
+                'digraph { Add [op_type=Add] output [op_type=Identity] input1:0 -> Add input2:0 -> ' 
                 'Add Add:0 -> output }',
                 onnx_to_graphviz(g))
 
@@ -200,9 +203,9 @@ def test_conv2d(self):
 
             g = process_tf_graph(sess.graph)
             self.assertEqual(
-                'digraph { Conv2D__2 [op_type=Transpose] kernel [op_type=Reshape] Conv2D__3 [op_type=Transpose] ' \
-                'Conv2D [op_type=Conv] Conv2D__4 [op_type=Transpose] output [op_type=Identity] input1:0 -> ' \
-                'Conv2D__2 k:0 -> kernel "kernel/shape":0 -> kernel kernel:0 -> Conv2D__3 Conv2D__2:0 -> Conv2D ' \
+                'digraph { Conv2D__2 [op_type=Transpose] kernel [op_type=Reshape] Conv2D__3 [op_type=Transpose] ' 
+                'Conv2D [op_type=Conv] Conv2D__4 [op_type=Transpose] output [op_type=Identity] input1:0 -> ' 
+                'Conv2D__2 k:0 -> kernel "kernel/shape":0 -> kernel kernel:0 -> Conv2D__3 Conv2D__2:0 -> Conv2D ' 
                 'Conv2D__3:0 -> Conv2D Conv2D:0 -> Conv2D__4 Conv2D__4:0 -> output }',
                 onnx_to_graphviz(g))
 
@@ -234,7 +237,7 @@ def test_reshape(self):
             _ = tf.identity(x_, name="output")
             g = process_tf_graph(sess.graph)
             self.assertEqual(
-                'digraph { Reshape [op_type=Reshape] output [op_type=Identity] input1:0 -> Reshape ' \
+                'digraph { Reshape [op_type=Reshape] output [op_type=Identity] input1:0 -> Reshape ' 
                 '"Reshape/shape":0 -> Reshape Reshape:0 -> output }',
                 onnx_to_graphviz(g))
 
@@ -257,7 +260,7 @@ def rewrite_test(g, ops):
             _ = tf.identity(x_, name="output")
             g = process_tf_graph(sess.graph, custom_rewriter=[rewrite_test])
             self.assertEqual(
-                'digraph { Add [op_type=Mul] output [op_type=Identity] input1:0 -> ' \
+                'digraph { Add [op_type=Mul] output [op_type=Identity] input1:0 -> '
                 'Add input1:0 -> Add Add:0 -> output }',
                 onnx_to_graphviz(g))
 

tests/test_internals.py

@@ -1,6 +1,9 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT license.
 
+from __future__ import division
+from __future__ import print_function
+
 import unittest
 from collections import namedtuple
 
@@ -136,7 +139,6 @@ def test_rewrite_subgraph(self):
                    'n3:0 -> ReplacedOp__2 ReplacedOp__2:0 -> n6 }'
         self.assertEqual(expected, result)
 
-
     def test_match_flipped(self):
         n1 = helper.make_node("Sub", ["i1", "i1"], ["n1:0"], name="n1")
         n2 = helper.make_node("Add", ["i2", "i2"], ["n2:0"], name="n2")

tf2onnx/__init__.py

@@ -1,12 +1,13 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT license.
 
-from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
 from .version import version as __version__
 
-__all__ = ["utils", "graph_matcher", "graph", "tf2onnx"]
+__all__ = ["utils", "graph_matcher", "graph", "tfonnx"]
+
+import tf2onnx
 from tf2onnx import tfonnx, utils, graph, graph_matcher

tf2onnx/convert.py

@@ -5,6 +5,9 @@
 python -m tf2onnx.convert : tool to convert a frozen tensorflow to onnx
 """
 
+from __future__ import division
+from __future__ import print_function
+
 import argparse
 import sys
 

tf2onnx/graph.py

@@ -4,8 +4,13 @@
 """
 tf2onnx.graph - class to manage graph manipulation on top of onnx
 """
+
+from __future__ import division
+from __future__ import print_function
+
+import tf2onnx
 from onnx import numpy_helper, optimizer, ModelProto, defs, OperatorSetIdProto
-from tf2onnx import utils, tfonnx, __version__
+from tf2onnx import utils, __version__
 from tf2onnx.utils import *
 
 
@@ -31,7 +36,7 @@ def __init__(self, node, graph):
         for a in node.attribute:
             self._attr[a.name] = a
         # try to find a dtype for this node
-        dtype = graph._dtypes.get(node.name)
+        dtype = graph.get_dtype(node.name)
         if not dtype:
             dtype = self._attr.get("dtype")
             if dtype:
@@ -230,7 +235,7 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
             dtypes: dict of tensorflow dtype
         """
         if target is None:
-            target = tfonnx.DEFAULT_TARGET
+            target = tf2onnx.tfonnx.DEFAULT_TARGET
         self._nodes = []
         self._initializers = {}
         self._nodes_by_name = {}

tf2onnx/tfonnx.py

@@ -4,6 +4,10 @@
 """
 tf2onnx.tf2onnx - rewrite tensorflow graph to onnx graph
 """
+
+from __future__ import division
+from __future__ import print_function
+
 import collections
 import logging
 import sys
@@ -87,7 +91,7 @@ def tensorflow_to_onnx(graph):
                 onnx_tensor = utils.tf_to_onnx_tensor(node.get_attr(a), name=node.name + ":0")
                 attr[a] = onnx_tensor
             elif a == "DstT":
-                attr["to"] =  utils.map_tf_dtype(node.get_attr("DstT"))
+                attr["to"] = utils.map_tf_dtype(node.get_attr("DstT"))
             elif a == "SrcT":
                 continue
             elif a in ignored_attr:
@@ -127,6 +131,7 @@ def _convert_shapenode_to_int64(ctx, node, input_number):
 
 # pylint: disable=W0613,C0111,W0612
 
+
 def no_op(ctx, node, name, args):
     """Skip node."""
     return None
@@ -304,12 +309,14 @@ def reshape_op5(ctx, node, name, args):
 HWCN_TO_NCHW = [3, 2, 0, 1]
 NCHW_TO_HWCN = [2, 3, 1, 0]
 
+
 def spatial_map(shape, perm):
     new_shape = shape[:]
     for i in perm:
         new_shape[i] = shape[perm[i]]
     return new_shape
 
+
 def conv_convert_inputs(ctx, node, with_kernel=False, new_kernel_shape=None,
                         input_indices=None, output_indices=None):
     """Convert input and kernel from tensorflow to onnx. This maybe require to
@@ -988,6 +995,7 @@ def onehot_op(ctx, node, name, args):
         return [node, transpose_op]
     return node
 
+
 def fused_batchnorm_op7(ctx, node, name, args):
     node.type = "BatchNormalization"
     # tf inputs: x, scale, bias, mean, variance
@@ -999,7 +1007,6 @@ def fused_batchnorm_op7(ctx, node, name, args):
     return nodes
 
 
-
 # pylint: enable=W0613,C0111,W0612
 
 # map tensorflow ops to onnx ops. The format below is
@@ -1132,6 +1139,7 @@ def fused_batchnorm_op7(ctx, node, name, args):
     (7, _OPSET_7),
 ]
 
+
 def rewrite_random_uniform(g, ops):
     pattern = \
         OpTypePattern('Add', name='output', inputs=[
@@ -1315,7 +1323,7 @@ def tensorflow_onnx_mapping(g, continue_on_error, custom_op_handlers):
 def tf_optimize(sess, inputs, outputs, graph_def):
     """Optimize tensorflow graph for inference."""
     transforms = [
-        #"fold_constants(ignore_errors=true)",
+        # "fold_constants(ignore_errors=true)",
         "fold_batch_norms",
         "fold_old_batch_norms",
     ]

tf2onnx/utils.py

@@ -5,6 +5,9 @@
 tf2onnx.utils - misc utilities for tf2onnx
 """
 
+from __future__ import division
+from __future__ import print_function
+
 import numpy as np
 import tensorflow as tf
 from onnx import helper, onnx_pb
@@ -149,6 +152,7 @@ def get_shape(node):
         pass
     return dims
 
+
 def map_tf_dtype(dtype):
     if dtype:
         dtype = TF_TO_ONNX_DTYPE[dtype]
@@ -157,8 +161,6 @@ def map_tf_dtype(dtype):
 
 def node_name(name):
     """Get node name without io#."""
-    # FIXME: do we use this ?
-    assert isinstance(name, str)
     pos = name.find(":")
     if pos >= 0:
         return name[:pos]