init

Author: tjingrant

__init__.py

@@ -0,0 +1,77 @@
+"""Backend for running ONNX on Tensorflow
+
+To run this, you will need to have Tensorflow installed as well.
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import collections
+
+import numpy as np
+from onnx import onnx_pb2, checker
+from onnx.onnx_pb2 import GraphProto, TensorProto, AttributeProto
+import onnx.numpy_helper
+import onnx.defs
+from onnx.backend.base import Backend, BackendRep, Device, DeviceType, namedtupledict
+
+import numpy as np
+from onnx import onnx_pb2, helper
+import tensorflow as tf
+
+
+def get_device_option(device):
+	m = {DeviceType.CPU: '/cpu:0',
+		 DeviceType.CUDA: '/gpu:0'}
+	return m[device.type]
+
+def get_type(type):
+	t = {
+		""
+	}
+
+class TensorflowBackend(Backend):
+	@classmethod
+	def run_node(cls, node, inputs, device='CPU'):
+		super(TensorflowBackend, cls).run_node(node, inputs, device)
+
+		device_option = get_device_option(Device(device))
+		input_tensors = []
+		for i in inputs:
+			input_tensors.append(tf.constant(i))
+
+		if isinstance(inputs, dict):
+			feed_dict_raw = inputs
+		else:
+			assert(len(node.input) == len(inputs))
+			feed_dict_raw = dict(zip(node.input, inputs))
+
+		input_dict = dict(map(lambda x: (x[0], tf.constant(x[1])), feed_dict_raw.items()))
+		print(input_dict)
+		outputs = cls._onnx_node_to_tensorflow_op(node, input_dict)
+		output_dict = {}
+		with tf.Session() as sess:
+			with tf.device(device_option):
+				for key, val in outputs.items():
+					output_dict[key] = sess.run(val)
+		return output_dict
+
+	@classmethod
+	def _onnx_node_to_tensorflow_op(cls, node, input_dict):
+		def _merge_two_dicts(x, y):
+			z = x.copy()
+			z.update(y)
+			return z
+		output_dict = dict()
+		method_to_call = getattr(cls, "handle_" + node.op_type.lower())
+		output_dict = _merge_two_dicts(output_dict, method_to_call(node, input_dict))
+		return output_dict
+
+	@classmethod
+	def handle_relu(cls, node, input_dict):
+		output_name = node.output[0]
+		input_name = node.input[0]
+		return dict([(output_name, tf.nn.relu(input_dict[input_name]))])
+
+run_node = TensorflowBackend.run_node

onnxtf/__init__.py

@@ -0,0 +1,86 @@
+from __future__ import print_function
+import re
+import pprint
+import urllib2
+
+def get_tf_defs():
+	tf_op_defs = {}
+
+	def clean(name, input):
+		return input.replace(name + "(\"", "").replace("\")", "").split(":")[0].strip();
+
+	url = "https://raw.githubusercontent.com/tensorflow/tensorflow/master/tensorflow/core/ops/{}.cc"
+	fnames = ["nn_ops", "math_ops", "array_ops"]
+
+	for fname in fnames:
+
+		content = urllib2.urlopen(url.format(fname)).read()
+		content = re.sub("\([\s\n]+\"", "(\"", content)
+		content = re.sub("\"[\s\n]+\"", "", content)
+		content = content.split("\n")
+		content = [x.strip() for x in content]
+
+		in_op_def = False
+		curr_op = {}
+		for line in content:
+			if (in_op_def):
+				if line.startswith(".Input"):
+					name = clean(".Input", line)
+					# print("\ti " + name)
+					curr_op["i"].append(name)
+				elif line.startswith(".Output"):
+					name = clean(".Output", line)
+					# print("\to " + name)
+					curr_op["o"].append(name)
+				elif line.startswith(".Attr"):
+					name = clean(".Attr", line)
+					# print("\ta " + name)
+					curr_op["a"].append(name)
+				else:
+					in_op_def = False
+					tf_op_defs[curr_op["n"]] = curr_op
+			else:
+				if (line.startswith("REGISTER_OP")):
+					in_op_def = True
+					name = clean("REGISTER_OP", line)
+					curr_op = {
+						"n": name,
+						"i": [],
+						"o": [],
+						"a": []
+					}
+					# print(name)
+	return tf_op_defs
+
+tf_op_defs = get_tf_defs()
+pp = pprint.PrettyPrinter(indent=2)
+# pp.pprint(tf_op_defs)
+
+from onnx import onnx_pb2, helper, defs
+node_def = helper.make_node("Relu", ["X"], ["Y"])
+print(node_def)
+
+all_schemas = defs.get_all_schemas()
+# print(all_schemas)
+# for name in all_schemas:
+# 	# print(dir(all_schemas[name]))
+# 	print("i", all_schemas[name].input_desc)
+# 	print("o", all_schemas[name].output_desc)
+# 	print("a", all_schemas[name].attributes)
+all_schemas = dict(filter(lambda x: len(x[1].input_desc)==1 ,all_schemas.items()))
+all_schemas = dict(filter(lambda x: len(x[1].attributes)==0 or len(x[1].attributes)==1,all_schemas.items()))
+
+intersection = (set(tf_op_defs.keys()).intersection(set(all_schemas.keys())))
+
+for key in intersection:
+	print(key)
+	pp.pprint(tf_op_defs[key])
+	print("i", all_schemas[key].input_desc)
+	print("o", all_schemas[key].output_desc)
+	print("a", all_schemas[key].attributes)
+
+
+# print(intersection)
+
+
+

onnxtf/convert.py

@@ -0,0 +1,11 @@
+from setuptools import setup
+
+setup(name='onnxtf',
+      version='0.1',
+      description='Tensorflow backend for ONNX',
+      url='TBD',
+      author='IBM',
+      author_email='TBD',
+      license='TBD',
+      packages=['onnxtf'],
+      zip_safe=False)

scaffold/extract_tf_ops.py

@@ -0,0 +1,20 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import unittest
+import numpy as np
+from onnxtf.convert import run_node
+from onnx import helper
+
+class TestStringMethods(unittest.TestCase):
+
+    def test_relu(self):
+    	node_def = helper.make_node("Relu", ["X"], ["Y"])
+    	input = np.random.uniform(-1,1,1000)
+    	output = run_node(node_def, [input])
+    	np.testing.assert_almost_equal(output["Y"], np.clip(input, 0, 1))
+
+if __name__ == '__main__':
+    unittest.main()