diff --git a/nnvm/python/nnvm/frontend/tensorflow.py b/nnvm/python/nnvm/frontend/tensorflow.py index ad7c4fc6796f..973aefdfab24 100644 --- a/nnvm/python/nnvm/frontend/tensorflow.py +++ b/nnvm/python/nnvm/frontend/tensorflow.py @@ -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( @@ -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]) @@ -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(), @@ -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. @@ -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 @@ -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] @@ -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 @@ -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) @@ -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)) @@ -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. @@ -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 diff --git a/nnvm/tests/python/frontend/tensorflow/test_forward.py b/nnvm/tests/python/frontend/tensorflow/test_forward.py index d73080d1cb00..900ef7436210 100644 --- a/nnvm/tests/python/frontend/tensorflow/test_forward.py +++ b/nnvm/tests/python/frontend/tensorflow/test_forward.py @@ -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 = {} @@ -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 @@ -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() @@ -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 # ------- @@ -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) @@ -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()