From d29243038b0f6071851e53e945597c069f4b2e98 Mon Sep 17 00:00:00 2001 From: "honghua.cao" Date: Fri, 28 Aug 2020 23:23:33 +0800 Subject: [PATCH 1/2] iadd conv2d_transpose alter layout --- python/tvm/relay/op/nn/_nn.py | 43 ++++++++++++++ .../relay/test_pass_convert_op_layout.py | 58 ++++++++++++++++++- 2 files changed, 100 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index 2f0966c62b65..43fca6d5f80f 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -192,6 +192,49 @@ def legalize_conv2d_transpose(attrs, inputs, types): """ return topi.nn.conv2d_transpose_legalize(attrs, inputs, types) +@reg.register_convert_op_layout("nn.conv2d_transpose") +def convert_conv2d_transpose(attrs, inputs, tinfos, desired_layouts): + """Convert Layout pass registration for conv2d_transpose op. + + Parameters + ---------- + attrs : tvm.ir.Attrs + Attributes of current convolution + inputs : list of tvm.relay.Expr + The args of the Relay expr to be legalized + tinfos : list of types + List of input and output types + desired_layouts : list of layout strings + List of layouts defining our desired + layout for the data and kernel inputs respectively. + + Returns + ------- + result : tvm.relay.Expr + The transformed expr + """ + # pylint: disable=import-outside-toplevel + from tvm import relay + data, weight = inputs + new_attrs = dict(attrs) + assert len(desired_layouts) == 2, "A desired layout is expected for both of nn.conv2d's inputs" + desired_data_layout, desired_kernel_layout = map(str, desired_layouts) + assert desired_data_layout != "default", "Data layout cannot be default" + new_attrs['data_layout'] = desired_data_layout + + if desired_kernel_layout != "default": + new_attrs['kernel_layout'] = desired_kernel_layout + return relay.nn.conv2d_transpose(data, weight, **new_attrs) + + # Handle default kernel layouts + if desired_data_layout == 'NCHW': + new_attrs['kernel_layout'] = 'OIHW' + return relay.nn.conv2d_transpose(data, weight, **new_attrs) + elif desired_data_layout == 'NHWC': + new_attrs['kernel_layout'] = 'HWIO' + return relay.nn.conv2d_transpose(data, weight, **new_attrs) + + raise ValueError("Layout %s is not yet supported." % desired_data_layout) # conv3d_transpose reg.register_strategy("nn.conv3d_transpose", strategy.conv3d_transpose_strategy) diff --git a/tests/python/relay/test_pass_convert_op_layout.py b/tests/python/relay/test_pass_convert_op_layout.py index f3cdbfc86e51..aec758d7b5fe 100644 --- a/tests/python/relay/test_pass_convert_op_layout.py +++ b/tests/python/relay/test_pass_convert_op_layout.py @@ -90,6 +90,41 @@ def expected(): assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) +def test_conv_transpose_convert_layout(): + def before(): + x = relay.var("x", shape=(1, 56, 56, 64)) + weight = relay.var('weight', shape=(3, 3, 64, 64)) + y = relay.nn.conv2d_transpose(x, weight, + channels=64, + kernel_size=(3, 3), + padding=(1, 1), + data_layout='NHWC', + kernel_layout='HWIO') + y = relay.nn.relu(y) + y = relay.Function([x, weight], y) + return y + + def expected(): + x = relay.var("x", shape=(1, 56, 56, 64)) + weight = relay.var('weight', shape=(3, 3, 64, 64)) + x = relay.layout_transform(x, 'NHWC', 'NCHW') + weight = relay.layout_transform(weight, 'HWIO', 'OIHW') + y = relay.nn.conv2d_transpose(x, weight, + channels=64, + kernel_size=(3, 3), + padding=(1, 1)) + y = relay.nn.relu(y) + y = relay.layout_transform(y, 'NCHW', 'NHWC') + y = relay.Function(relay.analysis.free_vars(y), y) + return y + + a = before() + a = run_opt_pass(a, transform.ConvertLayout({'nn.conv2d_transpose': ['NCHW', 'OIHW']})) + b = run_opt_pass(expected(), transform.InferType()) + + assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) + + def test_conv_bias_pool_convert_layout(): def before(): x = relay.var("x", shape=(1, 56, 56, 64)) @@ -680,6 +715,7 @@ def before(): x = relay.var("x", shape=(1, 64, 56, 56)) weight1 = relay.var("weight1", shape=(64, 3, 3, 64)) weight2 = relay.var("weight2", shape=(64, 3, 3, 64), dtype='int8') + weight3 = relay.var("weight3", shape=(64, 3, 3, 64)) out = relay.nn.conv2d(x, weight1, channels=64, kernel_size=(3, 3), @@ -697,6 +733,13 @@ def before(): padding=(1, 1), data_layout='NCHW', kernel_layout='OHWI') + out = relay.cast(out, 'float32') + out = relay.nn.conv2d_transpose(out, weight3, + channels=64, + kernel_size=(3, 3), + padding=(1, 1), + data_layout='NCHW', + kernel_layout='OHWI') out = relay.Function(analysis.free_vars(out), out) return out @@ -704,6 +747,7 @@ def expected(): x = relay.var("x", shape=(1, 64, 56, 56)) weight1 = relay.var("weight1", shape=(64, 3, 3, 64)) weight2 = relay.var("weight2", shape=(64, 3, 3, 64), dtype='int8') + weight3 = relay.var("weight3", shape=(64, 3, 3, 64)) x = relay.layout_transform(x, 'NCHW', 'NHWC') weight1 = relay.layout_transform(weight1, 'OHWI', 'HWIO') out = relay.nn.conv2d(x, weight1, @@ -725,12 +769,23 @@ def expected(): padding=(1, 1), data_layout='NCHW', kernel_layout='OIHW') + out = relay.cast(out, 'float32') + out = relay.layout_transform(out, 'NCHW', 'NHWC') + weight3 = relay.layout_transform(weight3, 'OHWI', 'HWIO') + out = relay.nn.conv2d_transpose(out, weight3, + channels=64, + kernel_size=(3, 3), + padding=(1, 1), + data_layout='NHWC', + kernel_layout='HWIO') + out = relay.layout_transform(out, 'NHWC', 'NCHW') out = relay.Function(analysis.free_vars(out), out) return out a = before() desired_layouts = {'nn.conv2d': ['NHWC', 'HWIO'], - 'qnn.conv2d': ['NCHW', 'OIHW']} + 'qnn.conv2d': ['NCHW', 'OIHW'], + 'nn.conv2d_transpose': ['NHWC', 'HWIO'],} a = run_opt_pass(a, transform.ConvertLayout(desired_layouts)) b = run_opt_pass(expected(), transform.InferType()) @@ -751,5 +806,6 @@ def expected(): test_qnn_conv_concat_convert_layout() test_qnn_conv_add_convert_layout() test_conv_convert_kernel_layout() + test_conv_transpose_convert_layout() test_default_keyword() test_different_ops_convert_layout() From 32dfb00356017e62e555ea3a50790150de87f17a Mon Sep 17 00:00:00 2001 From: "honghua.cao" Date: Thu, 10 Sep 2020 17:38:09 +0800 Subject: [PATCH 2/2] [RELAY][OP] roi_align operator alter layout --- python/tvm/relay/op/vision/_rcnn.py | 38 +++++++++++++++++ python/tvm/relay/op/vision/rcnn.py | 5 ++- src/relay/op/vision/rcnn_op.cc | 41 +++++++++++++++++-- .../relay/test_pass_convert_op_layout.py | 31 ++++++++++++++ 4 files changed, 109 insertions(+), 6 deletions(-) diff --git a/python/tvm/relay/op/vision/_rcnn.py b/python/tvm/relay/op/vision/_rcnn.py index 6f5097df49d2..24d7de517494 100644 --- a/python/tvm/relay/op/vision/_rcnn.py +++ b/python/tvm/relay/op/vision/_rcnn.py @@ -26,6 +26,44 @@ reg.register_strategy("vision.roi_align", strategy.roi_align_strategy) reg.register_pattern("vision.roi_align", OpPattern.OUT_ELEMWISE_FUSABLE) +@reg.register_convert_op_layout("vision.roi_align") +def convert_roi_align(attrs, inputs, tinfos, desired_layouts): + """Convert Layout pass registration for roi_align op. + + Parameters + ---------- + attrs : tvm.ir.Attrs + Attributes of current roi_align + inputs : list of tvm.relay.Expr + The args of the Relay expr to be legalized + tinfos : list of types + List of input and output types + desired_layouts : list of layout strings + List of layouts defining our desired + layout for the data and rois inputs respectively. + + Returns + ------- + result : tvm.relay.Expr + The transformed expr + """ + # pylint: disable=import-outside-toplevel + from tvm import relay + data, rois = inputs + new_attrs = dict(attrs) + assert len(desired_layouts) == 2, \ + "A desired layout is expected for both of vision.roi_align's inputs" + desired_data_layout, desired_rois_layout = map(str, desired_layouts) + assert desired_data_layout != "default", "Data layout cannot be default" + assert desired_rois_layout == "default", "Rois layout must be default" + + new_attrs['layout'] = desired_data_layout + # rois layout not change + if desired_data_layout in ['NCHW', 'NHWC']: + return relay.vision.roi_align(data, rois, **new_attrs) + + raise ValueError("Layout %s is not yet supported." % desired_data_layout) + # roi_pool @reg.register_compute("vision.roi_pool") def compute_roi_pool(attrs, inputs, _): diff --git a/python/tvm/relay/op/vision/rcnn.py b/python/tvm/relay/op/vision/rcnn.py index 1798ae946dc0..f4edf91c2ace 100644 --- a/python/tvm/relay/op/vision/rcnn.py +++ b/python/tvm/relay/op/vision/rcnn.py @@ -24,7 +24,7 @@ def roi_align(data, rois, pooled_size, spatial_scale, sample_ratio=-1, layout='N Parameters ---------- data : relay.Expr - 4-D tensor with shape [batch, channel, height, width] + 4-D tensor with shape [batch, channel, height, width] or [batch, height, width, channel] rois : relay.Expr 2-D tensor with shape [num_roi, 5]. The last dimension should be in format of @@ -43,7 +43,8 @@ def roi_align(data, rois, pooled_size, spatial_scale, sample_ratio=-1, layout='N Returns ------- output : relay.Expr - 4-D tensor with shape [num_roi, channel, pooled_size, pooled_size] + 4-D tensor with shape [num_roi, channel, pooled_size, pooled_size] or + [num_roi, pooled_size, pooled_size, channel] """ return _make.roi_align(data, rois, pooled_size, spatial_scale, sample_ratio, layout) diff --git a/src/relay/op/vision/rcnn_op.cc b/src/relay/op/vision/rcnn_op.cc index f7e1ecb82dcb..382e38d2189f 100644 --- a/src/relay/op/vision/rcnn_op.cc +++ b/src/relay/op/vision/rcnn_op.cc @@ -25,6 +25,8 @@ #include #include +#include "../../transforms/infer_layout_util.h" + namespace tvm { namespace relay { @@ -43,14 +45,43 @@ bool ROIAlignRel(const Array& types, int num_inputs, const Attrs& attrs, CHECK(roi_align_attrs); CHECK_EQ(dshape.size(), 4) << "Input data should be 4-D."; CHECK_EQ(rshape.size(), 2) << "Input rois should be 2-D."; - CHECK_EQ(roi_align_attrs->layout, "NCHW") << "ROI Align only supports NCHW layout"; // assign output type - std::vector oshape( - {rshape[0], dshape[1], roi_align_attrs->pooled_size[0], roi_align_attrs->pooled_size[1]}); + std::vector oshape; + if (roi_align_attrs->layout == "NCHW") { + oshape = {rshape[0], dshape[1], roi_align_attrs->pooled_size[0], + roi_align_attrs->pooled_size[1]}; + } else { + CHECK_EQ(roi_align_attrs->layout, "NHWC") << "Unexpected ROI Align layout"; + oshape = {rshape[0], roi_align_attrs->pooled_size[0], roi_align_attrs->pooled_size[1], + dshape[3]}; + } + reporter->Assign(types[2], TensorType(oshape, data->dtype)); return true; } +template +Array > ROIAlignInferCorrectLayout(const Attrs& attrs, + const Array& new_in_layouts, + const Array& old_in_layouts, + const Array& old_in_types) { + // NOTE: Discard "const" qualifier here. + T* params = const_cast(attrs.as()); + + if (new_in_layouts.defined()) { + // Set the roi_align with the new layout. + CHECK_EQ(new_in_layouts.size(), 2); + } + + // We always make other operators to fit the layouts of roi_align layers and the roi_align + // transpose followed by layout_transpose, So this inference ignores all inputs. + + // Layout inference needs to define the layout for all inputs and output data layouts. + // For roi_align, the second inputs is 2-D tensor with shape [num_roi, 5]. + // So, we set the layouts as "N5". + return Array >{{params->layout, Layout("N5")}, {params->layout}}; +} + Expr MakeROIAlign(Expr data, Expr rois, Array pooled_size, double spatial_scale, int sample_ratio, String layout) { auto attrs = make_object(); @@ -78,7 +109,9 @@ RELAY_REGISTER_OP("vision.roi_align") .add_argument("data", "Tensor", "The input tensor.") .add_argument("rois", "Tensor", "The input rois") .set_support_level(5) - .add_type_rel("ROIAlign", ROIAlignRel); + .add_type_rel("ROIAlign", ROIAlignRel) + .set_attr("FInferCorrectLayout", + ROIAlignInferCorrectLayout); TVM_REGISTER_NODE_TYPE(ROIPoolAttrs); diff --git a/tests/python/relay/test_pass_convert_op_layout.py b/tests/python/relay/test_pass_convert_op_layout.py index e71cfdcf4ecc..09836d4581a5 100644 --- a/tests/python/relay/test_pass_convert_op_layout.py +++ b/tests/python/relay/test_pass_convert_op_layout.py @@ -713,6 +713,36 @@ def expected(): assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) +def test_roi_align_convert_layout(): + def before(): + x = relay.var("x", shape=(1, 64, 56, 56)) + rois = relay.var('rois', shape=(32, 5)) + y = relay.vision.roi_align(x, rois, + pooled_size=(14, 14), + spatial_scale=0.0625, + sample_ratio=2, + layout='NCHW') + y = relay.Function([x, rois], y) + return y + + def expected(): + x = relay.var("x", shape=(1, 64, 56, 56)) + rois = relay.var('rois', shape=(32, 5)) + x = relay.layout_transform(x, 'NCHW', 'NHWC') + y = relay.vision.roi_align(x, rois, + pooled_size=(14, 14), + spatial_scale=0.0625, + sample_ratio=2, + layout='NHWC') + y = relay.layout_transform(y, 'NHWC', 'NCHW') + y = relay.Function(relay.analysis.free_vars(y), y) + return y + + a = before() + a = run_opt_pass(a, transform.ConvertLayout({'vision.roi_align': ['NHWC', 'default']})) + b = run_opt_pass(expected(), transform.InferType()) + + assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a) def test_default_keyword(): """ Check that the default keyword selects correct TVM default layout. """ @@ -845,5 +875,6 @@ def expected(): test_qnn_conv_add_convert_layout() test_conv_convert_kernel_layout() test_conv_transpose_convert_layout() + test_roi_align_convert_layout() test_default_keyword() test_different_ops_convert_layout()