def testTensorNames(self):
fdef = self._build_function_def()
g, tensor_name_map = function_def_to_graph.function_def_to_graph_def(
fdef)
# Verify that inputs of body nodes are correctly renamed.
# foo_1
self.assertSequenceEqual(g.node[3].input, ["x:0", "y:0", "z:0"])
# foo_2
self.assertSequenceEqual(g.node[5].input,
["foo_1:0", "foo_1:1", "list_output:1"])
# Verify that the `tensor_name_map` has the correct mapping.
self.assertDictEqual(
tensor_name_map, {
"x": "x:0",
"^x": "^x",
"y": "y:0",
"^y": "^y",
"z": "z:0",
"^z": "^z",
"foo_1:d:0": "foo_1:0",
"foo_1:e:0": "foo_1:1",
"^foo_1": "^foo_1",
"list_output:a:0": "list_output:0",
"list_output:a:1": "list_output:1",
"^list_output": "^list_output",
"foo_2:d:0": "foo_2:0",
"foo_2:e:0": "foo_2:1",
"^foo_2": "^foo_2",
})
def testTensorNames(self):
fdef = self._build_function_def()
g, tensor_name_map = function_def_to_graph.function_def_to_graph_def(fdef)
# Verify that inputs of body nodes are correctly renamed.
# foo_1
self.assertSequenceEqual(g.node[3].input, ["x:0", "y:0", "z:0"])
# foo_2
self.assertSequenceEqual(g.node[5].input,
["foo_1:0", "foo_1:1", "list_output:1"])
# Verify that the `tensor_name_map` has the correct mapping.
self.assertDictEqual(
tensor_name_map, {
"x": "x:0",
"^x": "^x",
"y": "y:0",
"^y": "^y",
"z": "z:0",
"^z": "^z",
"foo_1:d:0": "foo_1:0",
"foo_1:e:0": "foo_1:1",
"^foo_1": "^foo_1",
"list_output:a:0": "list_output:0",
"list_output:a:1": "list_output:1",
"^list_output": "^list_output",
"foo_2:d:0": "foo_2:0",
"foo_2:e:0": "foo_2:1",
"^foo_2": "^foo_2",
})
def testShapes(self):
fdef = self._build_function_def()
g, _ = function_def_to_graph.function_def_to_graph_def(
fdef,
input_shapes=[tensor_shape.scalar(),
tensor_shape.vector(5), None])
self.assertEqual("shape" in g.node[0].attr, True)
self.assertSequenceEqual(
tensor_shape.TensorShape(g.node[0].attr["shape"].shape).as_list(), [])
self.assertEqual(g.node[0].attr["shape"].shape.unknown_rank, False)
self.assertEqual("shape" in g.node[1].attr, True)
self.assertSequenceEqual(
tensor_shape.TensorShape(g.node[1].attr["shape"].shape).as_list(), [5])
self.assertEqual(g.node[0].attr["shape"].shape.unknown_rank, False)
self.assertFalse("shape" in g.node[2].attr)
def testAttributesForArgDef(self):
@function.defun
def fn(x):
return x
inp = constant_op.constant(1.0)
fdef = fn.get_concrete_function(inp).function_def
fdef.arg_attr[0].attr["_test_attr"].s = "value".encode("ascii")
graph_def = function_def_to_graph.function_def_to_graph_def(fdef)
placeholders = [
ndef for ndef in graph_def[0].node if ndef.op == "Placeholder"
]
self.assertEqual(1, len(placeholders))
self.assertEqual(placeholders[0].attr["_test_attr"].s,
"value".encode("ascii"))
def testShapes(self):
fdef = self._build_function_def()
g, _ = function_def_to_graph.function_def_to_graph_def(
fdef,
input_shapes=[tensor_shape.scalar(),
tensor_shape.vector(5), None])
self.assertEqual("shape" in g.node[0].attr, True)
self.assertSequenceEqual(
tensor_shape.TensorShape(g.node[0].attr["shape"].shape).as_list(), [])
self.assertEqual(g.node[0].attr["shape"].shape.unknown_rank, False)
self.assertEqual("shape" in g.node[1].attr, True)
self.assertSequenceEqual(
tensor_shape.TensorShape(g.node[1].attr["shape"].shape).as_list(), [5])
self.assertEqual(g.node[0].attr["shape"].shape.unknown_rank, False)
self.assertFalse("shape" in g.node[2].attr)
def testAttributesForArgDef(self):
@function.defun
def fn(x):
return x
inp = constant_op.constant(1.0)
fdef = fn.get_concrete_function(inp).function_def
fdef.arg_attr[0].attr["_test_attr"].s = "value".encode("ascii")
graph_def = function_def_to_graph.function_def_to_graph_def(fdef)
placeholders = [
ndef for ndef in graph_def[0].node if ndef.op == "Placeholder"
]
self.assertEqual(1, len(placeholders))
self.assertEqual(placeholders[0].attr["_test_attr"].s,
"value".encode("ascii"))
def _partition_call_operator(self, inputs, attr):
"""
Convert the Relay Partition call ops into Relay Function calls and
function definitions from Tensorflow graph library attribute to Relay global
functions
Parameters
----------
node: TensorFlow graph node object.
A TensorFlow graph node object.
inputs : List[tvm.relay.Expr]
List of input symbols.
attrs : Dict[tvm.Attrs]
Dict of operator attributes.
Returns
-------
op : tvm.relay.Expr
Converted relay expression.
"""
try:
from tensorflow.python.framework import function_def_to_graph
except ImportError as e:
raise ImportError(
"Unable to import tensorflow which is required {}".format(e))
main_graph_proto = self._main_graph_proto
outer_graph_def = main_graph_proto._graph
node_func_name = attr.get("f").name
func = next(
(f for f in outer_graph_def.library.function
if f.signature.name == node_func_name),
None,
)
if func:
devices = set(node.device for node in func.node_def)
if len(devices) > 1:
raise Exception("Found inconsistent Device assignment in the "
"Stateful Partitioned SubGraph. Rejecting "
"the subgraph ")
# Convert function definition to graph
func_input_shapes = func.attr["_input_shapes"].list.shape
subgraph, _ = function_def_to_graph.function_def_to_graph_def(
func, func_input_shapes)
# Computing subgraph's input shape dictionary
subgraph_shape_dict, input_expr_dict = {}, {}
for f_arg, input in zip(func.signature.input_arg, inputs):
input_expr_dict[f_arg.name] = input
subgraph_shape_dict[f_arg.name] = _infer_shape(
input, main_graph_proto._mod)
func_name = "func_{}".format(func.signature.name)
try:
global_func = main_graph_proto._mod[func_name]
sub_func = global_func
sub_params = main_graph_proto._params
except ValueError:
# Construct relay nodes from the subgraph
g1 = SubGraphProto(main_graph_proto)
sub_func, sub_params = g1.from_tensorflow(
subgraph, shape=subgraph_shape_dict)
main_graph_proto._params.update(sub_params)
func_expr = _function.Function(sub_func.params, sub_func.body)
global_func = tvm.relay.GlobalVar(func_name)
main_graph_proto._mod[global_func] = func_expr
main_graph_proto._mod = InferType()(main_graph_proto._mod)
param_exprs = []
for param_expr in sub_func.params:
# sub_params is subset of sub_func.params
param_name = param_expr.vid.name_hint
if param_name in input_expr_dict.keys():
param_exprs.append(input_expr_dict[param_name])
elif param_name in sub_params.keys():
param_exprs.append(param_expr)
else:
raise Exception(
"Input parameter {} not found".format(param_name))
sb = tvm.relay.scope_builder.ScopeBuilder()
loop_ret = global_func(*param_exprs)
sb.ret(loop_ret)
ret = sb.get()
else:
raise Exception("Function not found - {}".format(node_func_name))
return ret
def from_tensorflow(graph_def, layout="NHWC", shape=None, outputs=None):
"""convert tensorflow2.x graph into relay function.
Parameters
----------
graph_def : must be frozen graph (no variables allowed).
Placeholders are assumed to be inputs to the graph.
tensorflow/core/framework/graph.proto
message GraphDef {
repeated NodeDef node = 1;
FunctionDefLibrary library = 2;
}
tensorflow/core/framework/function.proto
message FunctionDef {
repeated NodeDef node_def = 3;
}
layout : str
The layout for the model.
shape : List[str, List[int]]
Input to the model. It is a key and shape vector mapping. Applies to placeholders.
outputs : List[str]
The list of output nodes. The last node is treated as the output if not
specified.
Returns
-------
mod : tvm.IRModule
The module that optimizations will be performed on.
params : dict of str to tvm.nd.NDArray
Dict of converted parameters stored in tvm.nd.NDArray format.
Examples
--------
"x+1" tf module where x has a shape of (2,2) is converted as follows:
mod : tvm.IRModule
def @func___inference_add_95(%x: Tensor[(2, 2), float32], %add/y: Tensor[(2, 2), float32])
-> Tensor[(2, 2), float32] {
add(%x, %add/y) /* Identity */ /* ty=Tensor[(2, 2), float32] */
}
def @main(%x1: Tensor[(2, 2), float32], %add/y1: Tensor[(2, 2), float32]) {
@func___inference_add_95(%x1, %add/y1) /* Identity */
}
params : dict of str to tvm.nd.NDArray
{'add/y': <tvm.nd.NDArray shape=(2, 2), cpu(0)>
"""
# Subgraph graph_defs are cached here to avoid a TF error when parsing after prelude init
graph_def_library = {}
for func in graph_def.library.function:
inshape = func.attr["_input_shapes"].list.shape
graph_def_library[func.signature.name], _ = function_def_to_graph.function_def_to_graph_def(
func, inshape
)
module = RelayModule()
g = GraphProto(module)
func, params = g.from_tensorflow(graph_def, layout, shape, outputs, gdef_lib=graph_def_library)
module.mod["main"] = func
module.params.update(params)
return module.mod, module.params
