diff --git a/nnvm/python/nnvm/frontend/tensorflow.py b/nnvm/python/nnvm/frontend/tensorflow.py index 17f08cc3e4a7..3ab997617a20 100644 --- a/nnvm/python/nnvm/frontend/tensorflow.py +++ b/nnvm/python/nnvm/frontend/tensorflow.py @@ -33,6 +33,8 @@ def __call__(self, inputs, attrs, *args): self._ignores.append('_input_shapes') self._ignores.append('T') self._ignores.append('use_cudnn_on_gpu') + self._ignores.append('_node_name') + self._ignores.append('is_training') return AttrConvert(self._op_name, self._transforms, self._excludes, self._disables, self._ignores, self._extras, self._custom_check)(inputs, attrs, *args) @@ -230,6 +232,85 @@ def _impl(inputs, attr, params): custom_check=_dimension_constraint())(inputs, attr) return _impl +def _depthwise_conv(): + def _impl(inputs, attr, params): + attr['data_format'] = attr['data_format'].decode("utf-8") + input_shapes = attr['_input_shapes'][inputs[0]] + + # Extract kernel shape from params + conv_param_weights = params[inputs[1].list_output_names()[0]] + + if attr['data_format'] == 'NHWC': + kernel_h, kernel_w, _, depth_mult = conv_param_weights.shape + attr['kernel_shape'] = (conv_param_weights.shape[0], conv_param_weights.shape[1]) + attr['channels'] = input_shapes[0][3] * depth_mult + if 'dilations' in attr: + attr['dilations'] = (attr['dilations'][0], attr['dilations'][1]) + elif attr['data_format'] == 'NCHW': + depth_mult, _, kernel_h, kernel_w = conv_param_weights.shape + attr['kernel_shape'] = (conv_param_weights.shape[2], conv_param_weights.shape[3]) + attr['channels'] = input_shapes[0][1] * depth_mult + if 'dilations' in attr: + attr['dilations'] = (attr['dilations'][2], attr['dilations'][3]) + else: + raise TypeError("Unsupported data format type : {}".format(attr['data_format'])) + + # Fix strides + attr['strides'] = (attr['strides'][1], attr['strides'][2]) + + # Fix groups + attr['groups'] = attr['channels'] + + # Fix padding + attr['padding'] = attr['padding'].decode("utf-8") + + if attr['padding'] == 'VALID': + attr['padding'] = [0, 0] + elif attr['padding'] == 'SAME': + stride_h, stride_w = attr['strides'] + kernel_h, kernel_w = attr['kernel_shape'] + if attr['data_format'] == 'NHWC': + in_h = input_shapes[0][1] + in_w = input_shapes[0][2] + else: + in_h = input_shapes[0][2] + in_w = input_shapes[0][3] + + pad_v = _get_pad_pair(in_h, kernel_h, stride_h) + pad_h = _get_pad_pair(in_w, kernel_w, stride_w) + + if attr['data_format'] == 'NHWC': + inputs[0] = _sym.pad(data=inputs[0], + pad_width=((0, 0), + (pad_v[0], pad_v[1]), + (pad_h[0], pad_h[1]), + (0, 0))) + else: + inputs[0] = _sym.pad(data=inputs[0], + pad_width=((0, 0), + (0, 0), + (pad_v[0], pad_v[1]), + (pad_h[0], pad_h[1]))) + + attr['padding'] = [0, 0] + + else: + raise TypeError("Unsupported padding type : {}".format(attr['padding'])) + + if 'kernel_layout' not in attr: + attr['kernel_layout'] = 'HWOI' if attr['data_format'] == 'NHWC' else 'OIHW' + + return AttrCvt( + op_name=_dimension_picker('conv'), + transforms={ + 'kernel_shape': 'kernel_size', + 'data_format': 'layout', + 'dilations': ('dilation', (0, 0)), + 'group': ('groups', 1)}, + extras={'use_bias': len(inputs) == 3}, + custom_check=_dimension_constraint())(inputs, attr) + return _impl + def _decode_image(): def _impl(inputs, attr, params): # Image decode wrapper: Expecting user to feed decoded input to next layer drop this layer. @@ -358,9 +439,27 @@ def _impl(inputs, attr, params): op_name='batch_norm', transforms={'scale_after_normalization':'scale', 'variance_epsilon':'epsilon'}, extras={'axis': 3}, # Fix axis + ignores=['data_format'], disables=['momentum'])(new_inputs, attr) return _impl +def _relu6(): + def _impl(inputs, attr, params): + return _sym.clip(inputs[0], a_min=0, a_max=6) + return _impl + +def _shape(): + def _impl(inputs, attr, params): + input_shapes = attr['_input_shapes'][inputs[0]] + + # Fix the -1 dimensions to 1 + input_shapes[0] = [1 if x == -1 else x for x in input_shapes[0]] + params[attr['_node_name']] = tvm.nd.array(input_shapes[0]) + + return _sym.Variable(name=attr['_node_name'], + shape=params[attr['_node_name']].shape) + return _impl + # compatible operators that do NOT require any conversion. _identity_list = [] @@ -392,6 +491,10 @@ def _impl(inputs, attr, params): 'Add' : _elemwise('add'), 'Rsqrt' : _rsqrt(), 'Squeeze' : _squeeze(), + 'FusedBatchNorm' : _batch_norm(), + 'Relu6' : _relu6(), + 'DepthwiseConv2dNative' : _depthwise_conv(), + 'Shape' : _shape(), } @@ -458,9 +561,13 @@ def from_tensorflow(self, graph): self._num_input += 1 self._nodes[node.name] = _sym.Variable(name=node.name) - self._output_shapes[node.name] = \ - [tensor_util.TensorShapeProtoToList(shape) \ - for shape in self._parse_attr(node.attr)['_output_shapes']] + try: + self._output_shapes[node.name] = \ + [tensor_util.TensorShapeProtoToList(shape) \ + for shape in self._parse_attr(node.attr)['_output_shapes']] + except KeyError: + raise NotImplementedError( \ + "Please freeze the graph with add_shapes=True") elif node.op == "Const": # Assuming first Const node as Graph Input node if self._input_node == '': @@ -476,17 +583,29 @@ def from_tensorflow(self, graph): raise NotImplementedError( \ "Const {} couldn't be converted to Param.".format(node.name)) - self._output_shapes[node.name] = \ - [tensor_util.TensorShapeProtoToList(shape) \ - for shape in self._parse_attr(node.attr)['_output_shapes']] + try: + self._output_shapes[node.name] = \ + [tensor_util.TensorShapeProtoToList(shape) \ + for shape in self._parse_attr(node.attr)['_output_shapes']] + except KeyError: + raise NotImplementedError( \ + "Please freeze the graph with add_shapes=True") else: attr = self._parse_attr(node.attr) - self._output_shapes[node.name] = \ - [tensor_util.TensorShapeProtoToList(shape) for shape in attr['_output_shapes']] + try: + self._output_shapes[node.name] = \ + [tensor_util.TensorShapeProtoToList(shape) \ + for shape in attr['_output_shapes']] + except KeyError: + raise NotImplementedError( \ + "Please freeze the graph with add_shapes=True") # Pass the parsed shapes instead attr["_output_shapes"] = self._output_shapes[node.name] + # Pass the node name too in attr + attr["_node_name"] = node.name + try: inputs = [self._nodes[i] for i in node.input] input_shapes = {} diff --git a/nnvm/src/top/nn/convolution.cc b/nnvm/src/top/nn/convolution.cc index e6576baf893b..40c4d7a0e8c6 100644 --- a/nnvm/src/top/nn/convolution.cc +++ b/nnvm/src/top/nn/convolution.cc @@ -80,14 +80,7 @@ inline bool Conv2DInferShape(const nnvm::NodeAttrs& attrs, wshape = ConvertLayout(wshape, kOIHW, kernel_layout); - // Depthwise Expects weights in - // NCHW : CNHW(IOHW) - Conversion is handled in frontend. - // NHWC : HWCN(HWIO) - Original format - if (param.layout == "NHWC") { - wshape[kernel_layout.indexof('I')] *= param.groups; - } else { - wshape[0] *= param.groups; - } + wshape[kernel_layout.indexof('O')] *= param.groups; NNVM_ASSIGN_INPUT_SHAPE(attrs, *in_shape, Conv2DParam::kWeight, wshape); if (param.use_bias) { diff --git a/nnvm/tests/python/compiler/test_top_level2.py b/nnvm/tests/python/compiler/test_top_level2.py index b367d42e9a37..8f4c330d2efc 100644 --- a/nnvm/tests/python/compiler/test_top_level2.py +++ b/nnvm/tests/python/compiler/test_top_level2.py @@ -83,7 +83,7 @@ def test_grouped_conv2d_nchw(): def test_grouped_conv2d_nhwc(): x = sym.Variable("x") y = sym.conv2d(x, channels=32, kernel_size=(3,3), groups=32, - name="y", padding=(1,1), layout="NHWC", kernel_layout ='HWIO') + name="y", padding=(1,1), layout="NHWC", kernel_layout ='HWOI') dtype = "float32" dshape = (1, 18, 18, 32) kshape = (3, 3, 32, 1)