def _function_graph(TestClass):
f = TestClass().func
gdef = f.get_concrete_function().graph.as_graph_def()
gdef_ops = list(set([n.op for n in gdef.node]))
input_ = TestClass().get_input()
output = run_tf_code(f, input_)
return gdef, input_, output
def _model_graph(TestClass):
model = TestClass()
with tempfile.TemporaryDirectory() as model_path:
tf.saved_model.save(model, model_path)
imported = tf.saved_model.load(model_path)
f = imported.signatures["serving_default"]
gdef = f.graph.as_graph_def(add_shapes=True)
input_ = model.get_input()
output = run_tf_code(f, input_)
return gdef, input_, output
def model_graph(model, input_shape):
_input = get_input(input_shape)
f = save_and_reload(model(input_shape))
_output = run_tf_code(f, _input)
gdef = f.graph.as_graph_def(add_shapes=True)
return gdef, _input, _output