def test_compile_simple(self):
frozen_graph_model = _make_frozen_graph_model(lambda input_x, input_y, _: tf.add(input_x, input_y, name='z'))
compiled = compiler.compile_source(source=frozen_graph_model)
compiled_graph = compiled.model_proto.graph
self.assertEqual([graph_input.name for graph_input in compiled_graph.input], ['x:0', 'y:0'])
self.assertEqual(compiled.input_data_formats, [None, None])
def test_compile_dropout(self):
def _make_model(input_x, input_y, _):
return tf.nn.dropout(x=input_x + input_y, rate=0.5, name='z')
frozen_graph_model = _make_frozen_graph_model(_make_model)
compiled = compiler.compile_source(source=frozen_graph_model)
compiled_graph = compiled.model_proto.graph
self.assertEqual([graph_input.name for graph_input in compiled_graph.input], ['x:0', 'y:0'])
self.assertEqual(compiled.input_data_formats, [None, None])
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, name='w')
output_z = tf.multiply(input_x + input_y, weight, name='z')
session.run(weight.initializer)
return output_z
frozen_graph_model = _make_frozen_graph_model(_make_model)
compiled = compiler.compile_source(source=frozen_graph_model)
compiled_graph = compiled.model_proto.graph
self.assertEqual([graph_input.name for graph_input in compiled_graph.input], ['x:0', 'y:0'])
self.assertEqual(compiled.input_data_formats, [None, None])
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)
