Skip to content

[TensorFlow] Support for Shape and Reshape in TensorFlow frontend #1884

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ghost wants to merge 1 commit into from
Closed

[TensorFlow] Support for Shape and Reshape in TensorFlow frontend #1884

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 27 additions & 17 deletions nnvm/python/nnvm/frontend/tensorflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -387,12 +387,14 @@ def _reshape():
def _impl(inputs, attr, params):
try:
pop_node = inputs[1]
shape_arg = params.pop(pop_node.list_output_names()[0])
if pop_node.list_output_names()[0] in params.keys():
shape_arg = params.pop(pop_node.list_output_names()[0]).asnumpy()
elif '_output_shapes' in attr:
shape_arg = np.asarray(attr['_output_shapes'][0])
inputs.pop(1)

return AttrCvt(
op_name="reshape",
extras={'shape':tuple(shape_arg.asnumpy())},
extras={'shape':tuple(shape_arg)},
ignores=['Tshape'])(inputs, attr)
except KeyError:
return AttrCvt(
Expand Down Expand Up @@ -467,13 +469,6 @@ def _impl(inputs, attr, params):
return _sym.clip(inputs[0], a_min=0, a_max=6, name=attr['_node_name'])
return _impl

def _shape():
def _impl(inputs, attr, params):
# Result of this operator is prominently used by reshape operator.
# Just pass the input as it is so that reshape_like can be used there.
return inputs[0]
return _impl

def _fill():
def _impl(inputs, attr, params):
fill_arg = params.pop(inputs.pop(1).list_output_names()[0])
Expand Down Expand Up @@ -815,7 +810,6 @@ def _impl(inputs, attr, params):
'FusedBatchNormV2' : _fused_batch_norm(),
'Relu6' : _relu6(),
'DepthwiseConv2dNative' : _conv('depthwise'),
'Shape' : _shape(),
'Sigmoid' : AttrCvt('sigmoid'),
'Fill' : _fill(),
'GatherV2' : _gather_v2(),
Expand Down Expand Up @@ -1030,7 +1024,7 @@ def __init__(self):
self._num_param = 0
self._num_rnn_layer = False

def from_tensorflow(self, graph, layout="NHWC"):
def from_tensorflow(self, graph, layout="NHWC", shape_dict=None):
"""Construct nnvm nodes from tensorflow graph definition - GraphDef.

Follow the tensorflow graph definition to parse and convert it to NNVM.
Expand Down Expand Up @@ -1059,6 +1053,9 @@ def from_tensorflow(self, graph, layout="NHWC"):
params : dict
A dict of name: tvm.nd.array pairs, used as pretrained weights
"""
if shape_dict is not None:
self._params.update(shape_dict)
self._num_param = len(shape_dict)

try:
from tensorflow.python.framework import tensor_util
Expand Down Expand Up @@ -1111,6 +1108,19 @@ def from_tensorflow(self, graph, layout="NHWC"):

attr = self._parse_attr(node.attr)

elif node.op == "Shape":
inputs = [self._nodes[i] for i in node.input]
if node.input[0] in self._output_shapes:
input_shape = self._output_shapes[node.input[0]][0]
elif shape_dict is not None:
input_shape = _infer_out_shapes(inputs[0], params)[0]
else:
raise NotImplementedError("Shape not supported by NNVM")
self._params[node.name] = tvm.nd.array(np.asarray(input_shape))
self._num_param += 1
self._nodes[node.name] = _sym.Variable(name=node.name,
shape=np.asarray(input_shape).shape)

else:
# Pass the parsed shapes instead
attr["_output_shapes"] = self._output_shapes[node.name]
Expand Down Expand Up @@ -1141,7 +1151,6 @@ def from_tensorflow(self, graph, layout="NHWC"):
pass

inputs = self._fix_extranodes(node.op, attr, inputs)

op = self._convert_operator(node.op, inputs, attr, graph)
# Assuming only one output.
self._nodes[node.name] = op
Expand Down Expand Up @@ -1173,7 +1182,8 @@ def _parse_import_prerequisites(self, graph):
elif node.op == "Const":
pass
else:
if any([node.op in t for t in [_identity_list, _convert_map, _convert_map_rnn]]):
if any([node.op in t for t in [_identity_list, _convert_map,
_convert_map_rnn, 'Shape']]):
pass
else:
missing_operators.add(node.op)
Expand Down Expand Up @@ -1206,7 +1216,7 @@ def _parse_param(self, key, value, name):
self._nodes[name] = _sym.Variable(name=name,
shape=self._params[name].shape)
else:
if key != 'dtype' and key != '_output_shapes' and key != '_class':
if key not in ('dtype', '_output_shapes', '_class'):
raise NotImplementedError \
("Other attributes for a Const(param) Node {} ? .".format(key))

Expand Down Expand Up @@ -1350,7 +1360,7 @@ def _fix_extranodes(self, op_name, attr, inputs):

return inputs

def from_tensorflow(graph, layout="NHWC"):
def from_tensorflow(graph, layout="NHWC", shape_dict=None):
""" Load tensorflow graph which is a python tensorflow graph object into nnvm graph.
The companion parameters will be handled automatically.

Expand All @@ -1368,5 +1378,5 @@ def from_tensorflow(graph, layout="NHWC"):
Dict of converted parameters stored in tvm.ndarray format
"""
g = GraphProto()
sym, params = g.from_tensorflow(graph, layout)
sym, params = g.from_tensorflow(graph, layout, shape_dict)
return sym, params
32 changes: 27 additions & 5 deletions nnvm/tests/python/frontend/tensorflow/test_forward.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,14 +26,14 @@
#######################################################################
# Generic run functions for TVM & tensorflow
# ------------------------------------------
def run_tvm_graph(graph_def, input_data, input_node, num_output=1, target='llvm'):
def run_tvm_graph(graph_def, input_data, input_node, num_output=1, target='llvm', shape_dict=None):
""" Generic function to compile on nnvm and execute on tvm """

layout = None
if target == "cuda":
layout = "NCHW"

sym, params = nnvm.frontend.from_tensorflow(graph_def, layout=layout)
sym, params = nnvm.frontend.from_tensorflow(graph_def, layout=layout, shape_dict=shape_dict)
target_host = 'llvm'
if isinstance(input_data, list):
shape_dict = {}
Expand Down Expand Up @@ -88,7 +88,7 @@ def run_tf_graph(sess, input_data, input_node, output_node):
return output_data


def compare_tf_with_tvm(in_data, in_name, out_name, init_global_variables=False, no_gpu=False):
def compare_tf_with_tvm(in_data, in_name, out_name, init_global_variables=False, no_gpu=False, shape_dict=None):
"""Generic function to generate and compare tensorflow and TVM output"""

out_node = out_name.split(':')[0] if ":" in out_name else out_name
Expand Down Expand Up @@ -119,7 +119,7 @@ def compare_tf_with_tvm(in_data, in_name, out_name, init_global_variables=False,
if no_gpu and device == 'cuda':
continue

tvm_output = run_tvm_graph(final_graph_def, in_data, in_node, target=device)
tvm_output = run_tvm_graph(final_graph_def, in_data, in_node, target=device, shape_dict=None)
np.testing.assert_allclose(tf_output, tvm_output, atol=1e-5, rtol=1e-5)

sess.close()
Expand Down Expand Up @@ -259,6 +259,27 @@ def test_forward_reshape():
_test_reshape(np.arange(6), [3, -1])
_test_reshape(np.arange(6), [-1])

#######################################################################
# Shape
# -------

def _test_shape(data, data2):
""" One iteration of reshape operation with given data and out shape """

with tf.Graph().as_default():
in_data = array_ops.placeholder(shape=data.shape, dtype=data.dtype)
shape = array_ops.shape(data2)
array_ops.reshape(in_data, shape)
ishapes = {}
ishapes['Placeholder'] = data.shape
compare_tf_with_tvm(data, 'Placeholder:0', 'Reshape:0', shape_dict=ishapes)

def test_forward_shape():
_test_shape(np.arange(6.0), np.asarray([1,1,1,1,1,1]))
_test_shape(np.asarray([[10, 20]]), np.zeros((1,2)))
_test_shape(np.asarray([[[10, 20]]]), np.zeros([1,1,2]))

#######################################################################
#######################################################################
# Squeeze
# -------
Expand Down Expand Up @@ -937,7 +958,7 @@ def test_forward_leaky_relu():
with tf.Graph().as_default():
in1 = tf.placeholder(shape=inp_array.shape, dtype=inp_array.dtype)
tf.nn.leaky_relu(in1, alpha=0.4)
compare_tf_with_tvm(inp_array, 'Placeholder:0', 'LeakyRelu:0')
compare_tf_with_tvm(inp_array, 'Placeholder:0', 'LeakyRelu/mul:0')

def test_forward_elu():
ishape = (1, 3, 10, 10)
Expand Down Expand Up @@ -1007,6 +1028,7 @@ def test_forward_rel_ops():
# Transforms
test_forward_transpose()
test_forward_reshape()
test_forward_shape()
test_forward_squeeze()
test_forward_pack()
test_forward_resize_bilinear()
Expand Down