From 2ea34dbd3a0e493c221c085974080ae9f63a8b10 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Tue, 24 Dec 2019 00:04:34 -0500 Subject: [PATCH] [Relay/Topi][Op] Added native DepthToSpace and SpaceToDepth Operators (#4566) * Added tvm function stencil for subpixel operations to topi. * Topi subpixel operators added and tested. * Added subpixel attrs. * Added depth_to_space relay attributes. * depth_to_space fully working. * Fixed NHWC shape bug. * SpaceToDepth in and all tests passing. * lint fixes. * Added string include * Fixed topi formatting. * Added DCR/CDR mode to depthtospace operator. --- include/tvm/relay/attrs/nn.h | 20 +++ python/tvm/relay/frontend/onnx.py | 50 +------- python/tvm/relay/frontend/tensorflow.py | 66 +--------- python/tvm/relay/op/nn/_nn.py | 22 ++++ python/tvm/relay/op/nn/nn.py | 50 ++++++++ python/tvm/relay/op/op_attrs.py | 5 + src/relay/op/nn/nn.cc | 118 ++++++++++++++++++ tests/python/relay/test_op_level5.py | 65 ++++++++++ topi/python/topi/nn/__init__.py | 2 + topi/python/topi/nn/depth_to_space.py | 88 +++++++++++++ topi/python/topi/nn/space_to_depth.py | 80 ++++++++++++ topi/python/topi/testing/__init__.py | 2 + topi/python/topi/testing/depth_to_space.py | 54 ++++++++ topi/python/topi/testing/space_to_depth.py | 49 ++++++++ topi/tests/python/test_topi_depth_to_space.py | 86 +++++++++++++ topi/tests/python/test_topi_space_to_depth.py | 85 +++++++++++++ 16 files changed, 732 insertions(+), 110 deletions(-) create mode 100644 topi/python/topi/nn/depth_to_space.py create mode 100644 topi/python/topi/nn/space_to_depth.py create mode 100644 topi/python/topi/testing/depth_to_space.py create mode 100644 topi/python/topi/testing/space_to_depth.py create mode 100644 topi/tests/python/test_topi_depth_to_space.py create mode 100644 topi/tests/python/test_topi_space_to_depth.py diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index 046e043d7978..4e061bd7d256 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -821,6 +821,26 @@ struct DeformableConv2DAttrs : public tvm::AttrsNode { } }; +/*! \brief Attributes used in subpixel operators */ +struct SubPixelAttrs : public tvm::AttrsNode { + int block_size; + std::string layout; + std::string mode; + + TVM_DECLARE_ATTRS(SubPixelAttrs, "relay.attrs.SubPixelAttrs") { + TVM_ATTR_FIELD(block_size) + .describe("The size of subpixel blocks to compose or decompose.") + .set_default(1); + TVM_ATTR_FIELD(layout).set_default("NCHW").describe( + "Dimension ordering of input data. Can be 'NCHW', 'NHWC', etc." + "'N', 'C', 'H', 'W' stands for batch, channel, height, and width" + "dimensions respectively."); + TVM_ATTR_FIELD(mode).set_default("DCR").describe( + "Indicates order in which channels are accessed. Must be one of" + "DCR or CDR."); + } +}; // struct SubPixelAttrs + } // namespace relay } // namespace tvm #endif // TVM_RELAY_ATTRS_NN_H_ diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index c6f3a2e8714b..c7764db729ee 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -540,34 +540,7 @@ def _impl_v11(cls, inputs, attr, params): block_size = int(attr['blocksize']) mode = attr.get("mode", "DCR") - - # handle NCHW layout - indata = infer_value_simulated(inputs[0], params) - in_n, in_c, in_h, in_w = indata.shape - - # reshape to proper output - new_c = int(in_c / (block_size * block_size)) - new_h = in_h * block_size - new_w = in_w * block_size - newshape = (in_n, new_c, new_h, new_w) - - if mode == "DCR": - # expand input to larger dimension. - expanded = _op.reshape(inputs[0], - newshape=(in_n, block_size, block_size, new_c, in_h, in_w)) - # reorder to expand spatial blocks. - transposed = _op.transpose(expanded, axes=(0, 3, 4, 1, 5, 2)) - - else: # CRD mode - # expand input to larger dimension. - expanded = _op.reshape(inputs[0], - newshape=(in_n, new_c, block_size, block_size, in_h, in_w)) - # reorder to expand spatial blocks. - transposed = _op.transpose(expanded, axes=(0, 1, 4, 2, 5, 3)) - - return AttrCvt(op_name="reshape", - extras={'newshape': newshape}, - ignores=['mode', 'blocksize'])([transposed], attr) + return _op.nn.depth_to_space(inputs[0], block_size, mode=mode) class SpaceToDepth(OnnxOpConverter): @@ -578,26 +551,7 @@ class SpaceToDepth(OnnxOpConverter): def _impl_v1(cls, inputs, attr, params): block_size = int(attr['blocksize']) - - # handle NCHW layout - indata = infer_value_simulated(inputs[0], params) - in_n, in_c, in_h, in_w = indata.shape - - # reshape to proper output - new_c = in_c * (block_size * block_size) - new_h = int(in_h / block_size) - new_w = int(in_w / block_size) - newshape = (in_n, new_c, new_h, new_w) - - # expand input to larger dimension. - expanded = _op.reshape(inputs[0], - newshape=(in_n, in_c, new_h, block_size, new_w, block_size)) - # reorder to expand spatial blocks. - transposed = _op.transpose(expanded, axes=(0, 3, 5, 1, 2, 4)) - - return AttrCvt(op_name="reshape", - extras={'newshape': newshape}, - ignores=['blocksize'])([transposed], attr) + return _op.nn.space_to_depth(inputs[0], block_size) class Concat(OnnxOpConverter): diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 8526f540515d..a07f0dd29828 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -728,76 +728,18 @@ def _impl(inputs, attr, params): def _depth_to_space(): def _impl(inputs, attr, params): - # Need to handle data layouts differently. - input_shape = attr['_input_shapes'][inputs[0]] block_size = int(attr['block_size']) - if attr['data_format'].decode("utf-8") == 'NHWC': - in_n, in_h, in_w, in_c = input_shape - new_c = int(in_c / (block_size * block_size)) - - # First expand input to larger dimension. - expanded = _op.reshape( - inputs[0], newshape=(in_n, in_h, in_w, block_size, block_size, new_c)) - # Now reorder to expand spatial blocks. - transposed = _op.transpose(expanded, axes=(0, 1, 3, 2, 4, 5)) - # Finally reshape to proper output. - new_h = in_h * block_size - new_w = in_w * block_size - newshape = (in_n, new_h, new_w, new_c) - - else: # Handle NCHW layout - in_n, in_c, in_h, in_w = input_shape - new_c = int(in_c / (block_size * block_size)) - - expanded = _op.reshape( - inputs[0], newshape=(in_n, block_size, block_size, new_c, in_h, in_w)) - transposed = _op.transpose(expanded, axes=(0, 3, 4, 1, 5, 2)) - new_h = in_h * block_size - new_w = in_w * block_size - newshape = (in_n, new_c, new_h, new_w) - - return AttrCvt( - op_name="reshape", - extras={'newshape': newshape}, - ignores=['data_format', 'block_size'])([transposed], attr) + layout = attr['data_format'].decode("utf-8") + return _op.nn.depth_to_space(inputs[0], block_size, layout) return _impl def _space_to_depth(): def _impl(inputs, attr, params): - # Need to handle data layouts differently. - input_shape = attr['_input_shapes'][inputs[0]] block_size = int(attr['block_size']) - if attr['data_format'].decode("utf-8") == 'NHWC': - in_n, in_h, in_w, in_c = input_shape - new_h = int(in_h / block_size) - new_w = int(in_w / block_size) - - # First expand input to larger dimension. - expanded = _op.reshape( - inputs[0], newshape=(in_n, new_h, block_size, new_w, block_size, in_c)) - # Now reorder to expand spatial blocks. - transposed = _op.transpose(expanded, axes=(0, 1, 3, 2, 4, 5)) - # Finally reshape to proper output. - new_c = in_c * block_size * block_size - newshape = (in_n, new_h, new_w, new_c) - - else: # Handle NCHW layout - in_n, in_c, in_h, in_w = input_shape - new_h = int(in_h / block_size) - new_w = int(in_w / block_size) - - expanded = _op.reshape( - inputs[0], newshape=(in_n, in_c, new_h, block_size, new_w, block_size)) - transposed = _op.transpose(expanded, axes=(0, 3, 5, 1, 2, 4)) - new_c = int(in_c * block_size * block_size) - newshape = (in_n, new_c, new_h, new_w) - - return AttrCvt( - op_name="reshape", - extras={'newshape': newshape}, - ignores=['data_format', 'block_size'])([transposed], attr) + layout = attr['data_format'].decode("utf-8") + return _op.nn.space_to_depth(inputs[0], block_size, layout) return _impl diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index 73bb2e29f365..ce47736bc15d 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -904,6 +904,28 @@ def compute_cross_entropy_with_logits(attrs, inputs, out_dtype, target): x, y = inputs return [-topi.sum(x * y) / x.shape[0]] + +@reg.register_compute("nn.depth_to_space") +def compute_depth_to_space(attrs, inputs, out_dtype, target): + block_size = attrs.block_size + layout = attrs.layout + mode = attrs.mode + return [topi.nn.depth_to_space(inputs[0], block_size, layout=layout, mode=mode)] + +reg.register_schedule("nn.depth_to_space", schedule_injective) +reg.register_pattern("nn.depth_to_space", OpPattern.INJECTIVE) + + +@reg.register_compute("nn.space_to_depth") +def compute_space_to_depth(attrs, inputs, out_dtype, target): + block_size = attrs.block_size + layout = attrs.layout + return [topi.nn.space_to_depth(inputs[0], block_size, layout=layout)] + +reg.register_schedule("nn.space_to_depth", schedule_injective) +reg.register_pattern("nn.space_to_depth", OpPattern.INJECTIVE) + + # shape func @script def _conv2d_NCHWc_shape_func(dshape, kshape, strides, padding, dilation, oc_bn): diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 326d72fabc38..5b6e1740e75a 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -2079,3 +2079,53 @@ def cross_entropy_with_logits(predictions, targets): The computed result. """ return _make.cross_entropy_with_logits(predictions, targets) + + +def depth_to_space(data, block_size, layout='NCHW', mode='DCR'): + """Convert channels into spatial blocks. + + Parameters + ---------- + data : tvm.relay.Expr + Input data with channels divisible by block_size**2 + + block_size : int + Size of blocks to convert channels into. + + layout : string + One of NCHW or NHWC, indicates channel axis. + + mode : string + One of DCR or CDR, indicates which order channels + are accessed in. + + Returns + ------- + result : tvm.relay.Expr + Tensor with shape [in_batch, in_channel / block_size * block_size, + in_height * block_size, in_width * block_size] + """ + return _make.depth_to_space(data, block_size, layout, mode) + + +def space_to_depth(data, block_size, layout='NCHW'): + """Convert spatial blocks into channels. + + Parameters + ---------- + data : tvm.relay.Expr + Input data with spatial dimensions divisible by block_size + + block_size : int + Size of blocks to decompose into channels. + + layout : string + One of NCHW or NHWC, indicates channel axis. + + Returns + ------- + result : tvm.relay.Expr + Tensor with shape [in_batch, in_channel * block_size * block_size, + in_height / block_size, in_width / block_size] + """ + return _make.space_to_depth(data, block_size, layout) diff --git a/python/tvm/relay/op/op_attrs.py b/python/tvm/relay/op/op_attrs.py index e9ddca9cde9c..a0ce40b90632 100644 --- a/python/tvm/relay/op/op_attrs.py +++ b/python/tvm/relay/op/op_attrs.py @@ -299,3 +299,8 @@ class BinaryDenseAttrs(Attrs): @register_relay_attr_node class Conv2DTransposeAttrs(Attrs): """Attributes used in Transposed Conv2D operators""" + + +@register_relay_attr_node +class SubPixelAttrs(Attrs): + """Attributes used in depth to space and space to depth operators""" diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index 051419cd338e..ac3848513856 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -31,6 +31,7 @@ #include #include #include +#include #include "../type_relations.h" #include "../../pass/alter_op_layout.h" #include "../op_common.h" @@ -959,6 +960,123 @@ Accept logits. .set_support_level(10) .add_type_rel("CrossEntropy", CrossEntropyRel); +// Depth to space and space to depth +TVM_REGISTER_NODE_TYPE(SubPixelAttrs); + +bool DepthToSpaceRel(const Array& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 2); + const auto* data = types[0].as(); + if (data == nullptr) return false; + + static const Layout kNCHW("NCHW"); + + const SubPixelAttrs* param = attrs.as(); + CHECK(param != nullptr); + const int block_size = param->block_size; + const Layout in_layout(param->layout); + auto layout_converter = BijectiveLayoutNode::make(in_layout, kNCHW); + CHECK(layout_converter.defined()) + << "DepthToSpace only support input layouts that are convertible from NCHW." + << " But got " << in_layout; + + auto oshape = layout_converter.ForwardShape(data->shape); + oshape.Set(1, indexdiv(oshape[1], (block_size * block_size))); + oshape.Set(2, oshape[2] * block_size); + oshape.Set(3, oshape[3] * block_size); + + // Assign output type + reporter->Assign(types[1], + TensorTypeNode::make(layout_converter.BackwardShape(oshape), data->dtype)); + + return true; +} + +// Positional relay function to create DepthToSpace operator +// used by frontend FFI +Expr MakeDepthToSpace(Expr data, int block_size, std::string layout, std::string mode) { + auto attrs = make_node(); + attrs->block_size = block_size; + attrs->layout = std::move(layout); + attrs->mode = std::move(mode); + static const Op& op = Op::Get("nn.depth_to_space"); + return CallNode::make(op, {data}, Attrs(attrs), {}); +} + +TVM_REGISTER_API("relay.op.nn._make.depth_to_space").set_body_typed(MakeDepthToSpace); + +RELAY_REGISTER_OP("nn.depth_to_space") + .describe(R"code(Rearrange input channels into spatial pixels. + +- **data**: data is a 4D array of shape + (batch, in_channels, in_height, in_width) for NCHW + +- **out**: Output is a 4D array of shape + (batch, in_channels / block_size * block_size, in_height * block_size, in_width * block_size) for NCHW. + +)code" TVM_ADD_FILELINE) + .set_attrs_type() + .set_num_inputs(1) + .add_argument("data", "Tensor", "The input tensor") + .set_support_level(5) + .add_type_rel("DepthToSpace", DepthToSpaceRel); + +bool SpaceToDepthRel(const Array& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 2); + const auto* data = types[0].as(); + if (data == nullptr) return false; + + static const Layout kNCHW("NCHW"); + + const SubPixelAttrs* param = attrs.as(); + CHECK(param != nullptr); + const int block_size = param->block_size; + const Layout in_layout(param->layout); + auto layout_converter = BijectiveLayoutNode::make(in_layout, kNCHW); + CHECK(layout_converter.defined()) + << "SpaceToDepth only support input layouts that are convertible from NCHW." + << " But got " << in_layout; + + auto oshape = layout_converter.ForwardShape(data->shape); + oshape.Set(1, oshape[1] * (block_size * block_size)); + oshape.Set(2, indexdiv(oshape[2], block_size)); + oshape.Set(3, indexdiv(oshape[3], block_size)); + + // Assign output type + reporter->Assign(types[1], + TensorTypeNode::make(layout_converter.BackwardShape(oshape), data->dtype)); + + return true; +} + +// Positional relay function to create SpaceToDepth operator +// used by frontend FFI +Expr MakeSpaceToDepth(Expr data, int block_size, std::string layout) { + auto attrs = make_node(); + attrs->block_size = block_size; + attrs->layout = std::move(layout); + static const Op& op = Op::Get("nn.space_to_depth"); + return CallNode::make(op, {data}, Attrs(attrs), {}); +} + +TVM_REGISTER_API("relay.op.nn._make.space_to_depth").set_body_typed(MakeSpaceToDepth); + +RELAY_REGISTER_OP("nn.space_to_depth") + .describe(R"code(Rearrange spatial pixels into new output channels. + +- **data**: data is a 4D array of shape + (batch, in_channels, in_height, in_width) for NCHW + +- **out**: Output is a 4D array of shape + (batch, in_channels * block_size * block_size, in_height / block_size, in_width / block_size) for NCHW. + +)code" TVM_ADD_FILELINE) + .set_attrs_type() + .set_num_inputs(1) + .add_argument("data", "Tensor", "The input tensor") + .set_support_level(5) + .add_type_rel("SpaceToDepth", SpaceToDepthRel); } // namespace relay } // namespace tvm diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index f7447465c3ac..84e9f55d67e7 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -573,6 +573,69 @@ def test_run(batch, in_channel, size, out_channel, deformable_groups, groups): test_run(2, 4, 16, 4, 4, 1) +def test_depth_to_space(): + def verify_depth_to_space(dshape, block_size, layout, mode): + if layout == "NHWC": + out_shape = [dshape[0], dshape[1] * block_size, dshape[2] * block_size, dshape[3] / (block_size * block_size)] + else: + out_shape = [dshape[0], dshape[1] / (block_size * block_size), dshape[2] * block_size, dshape[3] * block_size] + + x_data = np.random.uniform(size=dshape).astype("float32") + if layout == "NHWC": + x_data = np.transpose(x_data, axes=[0, 3, 1, 2]) + ref_res = topi.testing.depth_to_space_python(x_data, block_size, mode=mode) + if layout == "NHWC": + x_data = np.transpose(x_data, axes=[0, 2, 3, 1]) + ref_res = np.transpose(ref_res, axes=[0, 2, 3, 1]) + + x = relay.var("x", relay.TensorType(dshape, "float32")) + z = relay.nn.depth_to_space(x, block_size, layout, mode) + assert "block_size=" in z.astext() + zz = run_infer_type(z) + assert zz.checked_type == relay.TensorType(ref_res.shape, "float32") + func = relay.Function([x], z) + + for target, ctx in ctx_list(): + for kind in ["graph", "debug"]: + intrp = relay.create_executor(kind, ctx=ctx, target=target) + op_res = intrp.evaluate(func)(x_data) + tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-4) + for layout in ["NHWC", "NCHW"]: + for mode in ["DCR", "CDR"]: + verify_depth_to_space((1, 4, 4, 4), 2, layout, mode) + + +def test_space_to_depth(): + def verify_space_to_depth(dshape, block_size, layout): + if layout == "NHWC": + out_shape = [dshape[0], dshape[1] / block_size, dshape[2] / block_size, dshape[3] * (block_size * block_size)] + else: + out_shape = [dshape[0], dshape[1] * (block_size * block_size), dshape[2] / block_size, dshape[3] / block_size] + + x_data = np.random.uniform(size=dshape).astype("float32") + if layout == "NHWC": + x_data = np.transpose(x_data, axes=[0, 3, 1, 2]) + ref_res = topi.testing.space_to_depth_python(x_data, block_size) + if layout == "NHWC": + x_data = np.transpose(x_data, axes=[0, 2, 3, 1]) + ref_res = np.transpose(ref_res, axes=[0, 2, 3, 1]) + + x = relay.var("x", relay.TensorType(dshape, "float32")) + z = relay.nn.space_to_depth(x, block_size, layout) + assert "block_size=" in z.astext() + zz = run_infer_type(z) + assert zz.checked_type == relay.TensorType(ref_res.shape, "float32") + func = relay.Function([x], z) + + for target, ctx in ctx_list(): + for kind in ["graph", "debug"]: + intrp = relay.create_executor(kind, ctx=ctx, target=target) + op_res = intrp.evaluate(func)(x_data) + tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-4) + for layout in ["NHWC", "NCHW"]: + verify_space_to_depth((1, 4, 4, 4), 2, layout) + + if __name__ == "__main__": test_resize_infer_type() test_resize() @@ -586,3 +649,5 @@ def test_run(batch, in_channel, size, out_channel, deformable_groups, groups): test_yolo_reorg() test_non_max_suppression() test_deformable_conv2d() + test_depth_to_space() + test_space_to_depth() \ No newline at end of file diff --git a/topi/python/topi/nn/__init__.py b/topi/python/topi/nn/__init__.py index 3aa4a4e9dbb8..b805b7c64919 100644 --- a/topi/python/topi/nn/__init__.py +++ b/topi/python/topi/nn/__init__.py @@ -42,3 +42,5 @@ from .sparse import * from .pad import * from .fifo_buffer import * +from .depth_to_space import * +from .space_to_depth import * diff --git a/topi/python/topi/nn/depth_to_space.py b/topi/python/topi/nn/depth_to_space.py new file mode 100644 index 000000000000..d847c08daf27 --- /dev/null +++ b/topi/python/topi/nn/depth_to_space.py @@ -0,0 +1,88 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=invalid-name +"""TVM operator depth_to_space compute.""" +from __future__ import absolute_import +import tvm +from .. import tag + + +def depth_to_space(data, block_size, layout='NCHW', mode='DCR'): + """Perform depth to space transformation on the data + + Parameters + ---------- + data : tvm.Tensor + 4-D tensor in either NCHW or NHWC layout. + + block_size : int + Size of blocks to compose from channel dimension. + + layout : string + Either NCHW or NHWC, indicating data layout. + + mode : string + Either DCR or CDR, indicates how channels should be accessed. + In DCR, channels are interwoven in the Tensorflow style while + in CDR channels are accessed sequentially as in Pytorch. + + Returns + ------- + output : tvm.Tensor + Output of shape [N, C / block_size**2, H * block_size, W * block_size] + """ + if layout == 'NCHW': + in_n, in_c, in_h, in_w = data.shape + channel_factor = tvm.truncdiv(in_c, (block_size * block_size)) + output_shape = [in_n, channel_factor, + in_h * block_size, in_w * block_size] + elif layout == 'NHWC': + in_n, in_h, in_w, in_c = data.shape + channel_factor = tvm.truncdiv(in_c, (block_size * block_size)) + output_shape = [in_n, in_h * block_size, + in_w * block_size, channel_factor] + else: + raise ValueError("Only NCHW and NHWC layouts are currently supported.") + + def _get_indices(*indices): + if layout == 'NCHW': + n, c, y, x = indices + elif layout == 'NHWC': + n, y, x, c = indices + return n, c, y, x + + def _get_pixel(n, c, y, x): + block_x = tvm.truncdiv(x, block_size) + block_y = tvm.truncdiv(y, block_size) + idx_x = tvm.truncmod(x, block_size) + idx_y = tvm.truncmod(y, block_size) + if mode == "DCR": + channel_idx = channel_factor * ((block_size * idx_y) + idx_x) + c + else: + channel_idx = (c * block_size * block_size) + ((block_size * idx_y) + idx_x) + + if layout == 'NCHW': + output = data(n, channel_idx, block_y, block_x) + else: + output = data(n, block_y, block_x, channel_idx) + return output + + def _compute(*indices): + n, c, y, x = _get_indices(*indices) + return _get_pixel(n, c, y, x) + + return tvm.compute(output_shape, _compute, name='depth_to_space', tag=tag.INJECTIVE) diff --git a/topi/python/topi/nn/space_to_depth.py b/topi/python/topi/nn/space_to_depth.py new file mode 100644 index 000000000000..6ed7cd64a448 --- /dev/null +++ b/topi/python/topi/nn/space_to_depth.py @@ -0,0 +1,80 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=invalid-name +"""TVM operator space_to_depth compute.""" +from __future__ import absolute_import +import tvm +from .. import tag + + +def space_to_depth(data, block_size, layout='NCHW'): + """Perform space to depth transformation on the data + + Parameters + ---------- + data : tvm.Tensor + 4-D tensor in either NCHW or NHWC layout. + + block_size : int + Size of blocks to decompose into channel dimension. + + layout : string + Either NCHW or NHWC, indicating data layout. + + Returns + ------- + output : tvm.Tensor + Output of shape [N, C * block_size**2, H / block_size, W / block_size] + """ + + if layout == 'NCHW': + in_n, in_c, in_h, in_w = data.shape + output_shape = [in_n, in_c * block_size * block_size, + tvm.truncdiv(in_h, block_size), tvm.truncdiv(in_w, block_size)] + elif layout == 'NHWC': + in_n, in_h, in_w, in_c = data.shape + output_shape = [in_n, tvm.truncdiv(in_h, block_size), tvm.truncdiv( + in_w, block_size), in_c * block_size * block_size] + else: + raise ValueError("Only NCHW and NHWC layouts are currently supported.") + + def _get_indices(*indices): + if layout == 'NCHW': + n, c, y, x = indices + elif layout == 'NHWC': + n, y, x, c = indices + return n, c, y, x + + def _get_pixel(n, c, y, x): + block_offset = tvm.truncdiv(c, in_c) + channel_idx = tvm.truncmod(c, in_c) + x_idx = tvm.truncmod(block_offset, block_size) + y_idx = tvm.truncdiv(block_offset, block_size) + + if layout == 'NCHW': + output = data(n, channel_idx, y_idx + + (y * block_size), x_idx + (x * block_size)) + else: + output = data(n, y_idx + (y * block_size), x_idx + + (x * block_size), channel_idx) + return output + + def _compute(*indices): + n, c, y, x = _get_indices(*indices) + return _get_pixel(n, c, y, x) + + return tvm.compute(output_shape, _compute, name='space_to_depth', tag=tag.INJECTIVE) diff --git a/topi/python/topi/testing/__init__.py b/topi/python/topi/testing/__init__.py index f9acb73bf0cc..a3e2241d4967 100644 --- a/topi/python/topi/testing/__init__.py +++ b/topi/python/topi/testing/__init__.py @@ -46,3 +46,5 @@ from .pool3d_python import pool3d_ncdhw_python from .pool_grad_python import pool_grad_nchw from .one_hot import one_hot +from .depth_to_space import depth_to_space_python +from .space_to_depth import space_to_depth_python diff --git a/topi/python/topi/testing/depth_to_space.py b/topi/python/topi/testing/depth_to_space.py new file mode 100644 index 000000000000..f4a60bc12d87 --- /dev/null +++ b/topi/python/topi/testing/depth_to_space.py @@ -0,0 +1,54 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=invalid-name, line-too-long, unused-variable, too-many-locals +"""Depth to space in python""" +import numpy as np + + +def depth_to_space_python(data, block_size, mode='DCR'): + """Depth to Space operator in python for NCHW layout. + + Parameters + ---------- + data : np.ndarray + 4-D with shape [batch, in_channel, in_height, in_width] + + block_size : int + Size of blocks to convert channel pixels into. + + Returns + ------- + d2s_out : np.ndarray + 4-D with shape [batch, in_channel / (block_size * block_size), + out_height * block_size, out_width * block_size] + """ + in_n, in_c, in_h, in_w = data.shape + new_h = int(in_h * block_size) + new_w = int(in_h * block_size) + new_c = int(in_c / (block_size * block_size)) + + if mode == 'DCR': + expanded = np.reshape( + data, newshape=[in_n, block_size, block_size, new_c, in_h, in_w]) + transposed = np.transpose(expanded, axes=[0, 3, 4, 1, 5, 2]) + else: + expanded = np.reshape( + data, newshape=(in_n, new_c, block_size, block_size, in_h, in_w)) + transposed = np.transpose(expanded, axes=(0, 1, 4, 2, 5, 3)) + newshape = [in_n, new_c, new_h, new_w] + d2s_out = np.reshape(transposed, newshape=newshape) + return d2s_out diff --git a/topi/python/topi/testing/space_to_depth.py b/topi/python/topi/testing/space_to_depth.py new file mode 100644 index 000000000000..3a3b94177602 --- /dev/null +++ b/topi/python/topi/testing/space_to_depth.py @@ -0,0 +1,49 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# pylint: disable=invalid-name, line-too-long, unused-variable, too-many-locals +"""Space to depth in python""" +import numpy as np + + +def space_to_depth_python(data, block_size): + """Space to Depth operator in python for NCHW layout. + + Parameters + ---------- + data : np.ndarray + 4-D with shape [batch, in_channel, in_height, in_width] + + block_size : int + Size of spatial blocks to decompose into channels. + + Returns + ------- + d2s_out : np.ndarray + 4-D with shape [batch, in_channel * (block_size * block_size), + out_height / block_size, out_width / block_size] + """ + in_n, in_c, in_h, in_w = data.shape + new_h = int(in_h / block_size) + new_w = int(in_h / block_size) + new_c = int(in_c * (block_size * block_size)) + + expanded = np.reshape( + data, newshape=[in_n, in_c, new_h, block_size, new_w, block_size]) + transposed = np.transpose(expanded, axes=[0, 3, 5, 1, 2, 4]) + newshape = [in_n, new_c, new_h, new_w] + d2s_out = np.reshape(transposed, newshape=newshape) + return d2s_out diff --git a/topi/tests/python/test_topi_depth_to_space.py b/topi/tests/python/test_topi_depth_to_space.py new file mode 100644 index 000000000000..4e895cb5db55 --- /dev/null +++ b/topi/tests/python/test_topi_depth_to_space.py @@ -0,0 +1,86 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Test code for depth to space""" +import numpy as np +import tvm +import topi +import topi.testing + +from common import get_all_backend + + +def verify_depth_to_space(block_size, batch, in_channel, in_height, in_width, layout='NCHW', mode='DCR'): + out_channel = int(in_channel / (block_size * block_size)) + out_height = int(in_height * block_size) + out_width = int(in_width * block_size) + + if layout == 'NCHW': + in_shape = [batch, in_channel, in_height, in_width] + out_shape = [batch, out_channel, out_height, out_width] + elif layout == 'NHWC': + in_shape = [batch, in_height, in_width, in_channel] + out_shape = [batch, out_height, out_width, out_channel] + else: + raise NotImplementedError('Layout not supported {}'.format(layout)) + + A = tvm.placeholder(in_shape, name='A', dtype='float32') + dtype = A.dtype + a_np = np.random.uniform(size=in_shape).astype(dtype) + + B = topi.nn.depth_to_space(A, block_size=block_size, layout=layout, mode=mode) + if layout == 'NHWC': + a_np = np.transpose(a_np, axes=[0, 3, 1, 2]) + b_np = topi.testing.depth_to_space_python(a_np, block_size, mode=mode) + if layout == 'NHWC': + a_np = np.transpose(a_np, axes=[0, 2, 3, 1]) + b_np = np.transpose(b_np, axes=[0, 2, 3, 1]) + + def check_device(device): + ctx = tvm.context(device, 0) + if not ctx.exist: + print("Skip because %s is not enabled" % device) + return + print("Running on target: %s" % device) + with tvm.target.create(device): + s = topi.generic.schedule_injective(B) + a = tvm.nd.array(a_np, ctx) + b = tvm.nd.array(np.zeros(out_shape, dtype=dtype), ctx) + f = tvm.build(s, [A, B], device) + f(a, b) + tvm.testing.assert_allclose(b.asnumpy(), b_np, rtol=1e-3, atol=1e-3) + + for device in get_all_backend(): + check_device(device) + + +def test_depth_to_space(): + for layout in ['NCHW', 'NHWC']: + for mode in ['DCR', 'CDR']: + # Simplest possible case + verify_depth_to_space(2, 1, 4, 1, 1, layout=layout, mode=mode) + # Average input size + verify_depth_to_space(2, 1, 32, 32, 32, layout=layout, mode=mode) + # Large block size + verify_depth_to_space(8, 1, 256, 32, 32, layout=layout, mode=mode) + # Large batch size + verify_depth_to_space(4, 8, 32, 32, 32, layout=layout, mode=mode) + # Large input size + verify_depth_to_space(4, 8, 32, 128, 128, layout=layout, mode=mode) + + +if __name__ == "__main__": + test_depth_to_space() diff --git a/topi/tests/python/test_topi_space_to_depth.py b/topi/tests/python/test_topi_space_to_depth.py new file mode 100644 index 000000000000..b25cad194301 --- /dev/null +++ b/topi/tests/python/test_topi_space_to_depth.py @@ -0,0 +1,85 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Test code for space to depth""" +import numpy as np +import tvm +import topi +import topi.testing + +from common import get_all_backend + + +def verify_space_to_depth(block_size, batch, in_channel, in_height, in_width, layout='NCHW'): + out_channel = int(in_channel * (block_size * block_size)) + out_height = int(in_height / block_size) + out_width = int(in_width / block_size) + + if layout == 'NCHW': + in_shape = [batch, in_channel, in_height, in_width] + out_shape = [batch, out_channel, out_height, out_width] + elif layout == 'NHWC': + in_shape = [batch, in_height, in_width, in_channel] + out_shape = [batch, out_height, out_width, out_channel] + else: + raise NotImplementedError('Layout not supported {}'.format(layout)) + + A = tvm.placeholder(in_shape, name='A', dtype='float32') + dtype = A.dtype + a_np = np.random.uniform(size=in_shape).astype(dtype) + + B = topi.nn.space_to_depth(A, block_size=block_size, layout=layout) + if layout == 'NHWC': + a_np = np.transpose(a_np, axes=[0, 3, 1, 2]) + b_np = topi.testing.space_to_depth_python(a_np, block_size) + if layout == 'NHWC': + a_np = np.transpose(a_np, axes=[0, 2, 3, 1]) + b_np = np.transpose(b_np, axes=[0, 2, 3, 1]) + + def check_device(device): + ctx = tvm.context(device, 0) + if not ctx.exist: + print("Skip because %s is not enabled" % device) + return + print("Running on target: %s" % device) + with tvm.target.create(device): + s = topi.generic.schedule_injective(B) + a = tvm.nd.array(a_np, ctx) + b = tvm.nd.array(np.zeros(out_shape, dtype=dtype), ctx) + f = tvm.build(s, [A, B], device) + f(a, b) + tvm.testing.assert_allclose(b.asnumpy(), b_np, rtol=1e-3, atol=1e-3) + + for device in get_all_backend(): + check_device(device) + + +def test_space_to_depth(): + for layout in ['NCHW', 'NHWC']: + # Simplest possible case + verify_space_to_depth(2, 1, 1, 2, 2, layout=layout) + # Average input size + verify_space_to_depth(2, 1, 32, 32, 32, layout=layout) + # Large block size + verify_space_to_depth(8, 1, 32, 64, 64, layout=layout) + # Large batch size + verify_space_to_depth(4, 8, 32, 32, 32, layout=layout) + # Large input size + verify_space_to_depth(4, 8, 32, 128, 128, layout=layout) + + +if __name__ == "__main__": + test_space_to_depth()