def test_compile_simple(self): tensorflow_model = _make_tensorflow_model(lambda input_x, input_y, _: tf.add(input_x, input_y, name='z')) compiled = compiler.compile_source(source=tensorflow_model) self.assertEqual(compiled.inputs, [Input(name='x:0', data_format=None), Input(name='y:0', data_format=None)]) self.assertEqual(compiled.outputs, ['z:0'])
def _make_frozen_graph_model(func): with eager_context.graph_mode(), tf.Graph().as_default(), tf.compat.v1.Session().as_default() as session: input_x = tf.compat.v1.placeholder(dtype=tf.float32, shape=[3, 4], name='x') input_y = tf.compat.v1.placeholder(dtype=tf.float32, shape=[3, 4], name='y') output_z = func(input_x, input_y, session) return frozen_graph_compiler.compile_source( source=TensorFlowModel(inputs=[Input(tensor=input_x), Input(tensor=input_y)], outputs=[output_z], session=session) )
def _make_onnx_model(func, batch_size_1, batch_size_2): with tf.Graph().as_default(), tf.compat.v1.Session().as_default() as session: input_x = tf.compat.v1.placeholder(dtype=tf.float32, shape=[batch_size_1, 4], name='x') input_y = tf.compat.v1.placeholder(dtype=tf.float32, shape=[batch_size_2, 4], name='y') output_z = func(input_x, input_y, session) frozen_graph_model = frozen_graph_compiler.compile_source( source=TensorFlowModel(inputs=[TfInput(tensor=input_x), TfInput(tensor=input_y)], outputs=[output_z], session=session) ) return onnx_compiler.compile_source(frozen_graph_model)
def test_compile_with_variables(self): def _make_model(input_x, input_y, session): weight = tf.Variable(initial_value=4.2, dtype=tf.float32) output_z = tf.multiply(input_x + input_y, weight, name='z') session.run(weight.initializer) return output_z frozen_graph_model = _make_tensorflow_model(_make_model) compiled = compiler.compile_source(source=frozen_graph_model) self.assertEqual(compiled.inputs, [Input(name='x:0', data_format=None), Input(name='y:0', data_format=None)]) self.assertEqual(compiled.outputs, ['z:0'])
def _make_onnx_model(): with tf.Graph().as_default(), tf.compat.v1.Session().as_default( ) as session: input_x = tf.compat.v1.placeholder(dtype=tf.float32, shape=[3, 4], name='x') input_y = tf.compat.v1.placeholder(dtype=tf.float32, shape=[3, 4], name='y') weight = tf.Variable(initial_value=4.2, dtype=tf.float32) output_z = tf.multiply(input_x + input_y, weight, name='z') session.run(weight.initializer) frozen_graph_model = tf_model_compiler.compile_source( source=TensorFlowModel( inputs=[Input( tensor=input_x), Input(tensor=input_y)], outputs=[output_z], session=session)) return frozen_graph_compiler.compile_source(frozen_graph_model)