From 84a9ea250f33994bfbcccdfc49318c4048c84076 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Wed, 18 Dec 2019 16:15:46 -0800 Subject: [PATCH 01/11] Added tvm function stencil for subpixel operations to topi. --- topi/python/topi/nn/depth_to_space.py | 58 +++++++++++++++++++++++++++ topi/python/topi/nn/space_to_depth.py | 55 +++++++++++++++++++++++++ 2 files changed, 113 insertions(+) create mode 100644 topi/python/topi/nn/depth_to_space.py create mode 100644 topi/python/topi/nn/space_to_depth.py 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..bedd806204fc --- /dev/null +++ b/topi/python/topi/nn/depth_to_space.py @@ -0,0 +1,58 @@ +# 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'): + 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) + channel_idx = channel_factor * ((block_size * idx_y) + idx_x) + c + + if layout == 'NCHW': + return data(n, channel_idx, block_y, block_x) + else: + return data(n, block_y, block_x, channel_idx) + + 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..d8d8a23b40e8 --- /dev/null +++ b/topi/python/topi/nn/space_to_depth.py @@ -0,0 +1,55 @@ +# 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'): + 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': + return data(n, channel_idx, y_idx + (y * block_size), x_idx + (x * block_size)) + else: + return data(n, y_idx + (y * block_size), x_idx + (x * block_size), channel_idx) + + 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) From 154c52970ec9f5f8454e0e5cf15bd7c59530e789 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Fri, 20 Dec 2019 19:04:33 -0500 Subject: [PATCH 02/11] Topi subpixel operators added and tested. --- topi/python/topi/nn/__init__.py | 2 + topi/python/topi/testing/__init__.py | 2 + topi/python/topi/testing/depth_to_space.py | 49 +++++++++++ topi/python/topi/testing/space_to_depth.py | 49 +++++++++++ topi/tests/python/test_topi_depth_to_space.py | 84 +++++++++++++++++++ topi/tests/python/test_topi_space_to_depth.py | 84 +++++++++++++++++++ 6 files changed, 270 insertions(+) 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/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/testing/__init__.py b/topi/python/topi/testing/__init__.py index 43e9f19d880b..931bcdb5780c 100644 --- a/topi/python/topi/testing/__init__.py +++ b/topi/python/topi/testing/__init__.py @@ -45,3 +45,5 @@ from .sequence_mask_python import sequence_mask 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..c3eab104c225 --- /dev/null +++ b/topi/python/topi/testing/depth_to_space.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 +"""Depth to space in python""" +import numpy as np + + +def depth_to_space_python(data, block_size): + """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)) + + 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]) + 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..5a1a2dc9a4fe --- /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..c4694dd27293 --- /dev/null +++ b/topi/tests/python/test_topi_depth_to_space.py @@ -0,0 +1,84 @@ +# 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'): + 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) + 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) + 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']: + # Simplest possible case + verify_depth_to_space(2, 1, 4, 1, 1, layout=layout) + # Average input size + verify_depth_to_space(2, 1, 32, 32, 32, layout=layout) + # Large block size + verify_depth_to_space(8, 1, 256, 32, 32, layout=layout) + # Large batch size + verify_depth_to_space(4, 8, 32, 32, 32, layout=layout) + # Large input size + verify_depth_to_space(4, 8, 32, 128, 128, layout=layout) + + +if __name__ == "__main__": + test_depth_to_space() \ No newline at end of file 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..12e93133e495 --- /dev/null +++ b/topi/tests/python/test_topi_space_to_depth.py @@ -0,0 +1,84 @@ +# 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() \ No newline at end of file From 32d27c9e2e7b7f7d1978b46ac878163172e00adc Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Fri, 20 Dec 2019 22:24:51 -0500 Subject: [PATCH 03/11] Added subpixel attrs. --- include/tvm/relay/attrs/nn.h | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index 046e043d7978..67bb94f12850 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -821,6 +821,20 @@ struct DeformableConv2DAttrs : public tvm::AttrsNode { } }; +/*! \brief Attributes used in subpixel operators */ +struct SubPixelAttrs : public tvm::AttrsNode { + std::string data_layout; + + TVM_DECLARE_ATTRS(SubPixelAttrs, "relay.attrs.SubPixelAttrs") { + TVM_ATTR_FIELD(data_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."); + } +}; // struct SubPixelAttrs + } // namespace relay } // namespace tvm #endif // TVM_RELAY_ATTRS_NN_H_ From f317f0d3d695107be117d4f587796654f8422c81 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Fri, 20 Dec 2019 23:39:12 -0500 Subject: [PATCH 04/11] Added depth_to_space relay attributes. --- include/tvm/relay/attrs/nn.h | 8 +++-- src/relay/op/nn/nn.cc | 59 ++++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 2 deletions(-) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index 67bb94f12850..354349624f58 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -823,10 +823,14 @@ struct DeformableConv2DAttrs : public tvm::AttrsNode { /*! \brief Attributes used in subpixel operators */ struct SubPixelAttrs : public tvm::AttrsNode { - std::string data_layout; + int block_size; + std::string layout; TVM_DECLARE_ATTRS(SubPixelAttrs, "relay.attrs.SubPixelAttrs") { - TVM_ATTR_FIELD(data_layout) + 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." diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index 051419cd338e..dc12ec4b7706 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -959,6 +959,65 @@ 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(data->shape[1], (block_size * block_size))); + oshape.Set(2, data->shape[2] * block_size); + oshape.Set(3, data->shape[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 opeator +// used by frontend FFI +Expr MakeDepthToSpace(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.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") + .add_type_rel("DepthToSpace", DepthToSpaceRel) + .set_attr("TOpPattern", kInjective); } // namespace relay } // namespace tvm From 7f988fa43fcc1ca151d87123d9d7f73a40071132 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Sat, 21 Dec 2019 17:24:39 -0500 Subject: [PATCH 05/11] depth_to_space fully working. --- python/tvm/relay/op/nn/_nn.py | 10 +++++++++ python/tvm/relay/op/nn/nn.py | 23 ++++++++++++++++++++ python/tvm/relay/op/op_attrs.py | 5 +++++ src/relay/op/nn/nn.cc | 4 ++-- tests/python/relay/test_op_level5.py | 32 ++++++++++++++++++++++++++++ 5 files changed, 72 insertions(+), 2 deletions(-) diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index e1372ac76480..5b20bed82b19 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -881,6 +881,16 @@ 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 + return [topi.nn.depth_to_space(inputs[0], block_size, layout=layout)] + +reg.register_schedule("nn.depth_to_space", schedule_injective) +reg.register_pattern("nn.depth_to_space", 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 fda5027ee49e..06354687afc8 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -1985,3 +1985,26 @@ 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'): + """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. + + 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) diff --git a/python/tvm/relay/op/op_attrs.py b/python/tvm/relay/op/op_attrs.py index 35b2c053f8cf..44ce60315bda 100644 --- a/python/tvm/relay/op/op_attrs.py +++ b/python/tvm/relay/op/op_attrs.py @@ -289,3 +289,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""" \ No newline at end of file diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index dc12ec4b7706..f126d477cba6 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -1016,8 +1016,8 @@ RELAY_REGISTER_OP("nn.depth_to_space") .set_attrs_type() .set_num_inputs(1) .add_argument("data", "Tensor", "The input tensor") - .add_type_rel("DepthToSpace", DepthToSpaceRel) - .set_attr("TOpPattern", kInjective); + .set_support_level(5) + .add_type_rel("DepthToSpace", DepthToSpaceRel); } // namespace relay } // namespace tvm diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index f7447465c3ac..3910a20be092 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -573,6 +573,37 @@ 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): + 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) + 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) + 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_depth_to_space((1, 4, 4, 4), 2, layout) + + if __name__ == "__main__": test_resize_infer_type() test_resize() @@ -586,3 +617,4 @@ 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() From 9d360dc304546b49b4bb3a3870073a2bd2752c14 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Sat, 21 Dec 2019 16:55:35 -0800 Subject: [PATCH 06/11] Fixed NHWC shape bug. --- python/tvm/relay/frontend/tensorflow.py | 33 ++----------------------- src/relay/op/nn/nn.cc | 6 ++--- 2 files changed, 5 insertions(+), 34 deletions(-) diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 460a14699a77..27a4d3c8d33e 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -664,38 +664,9 @@ 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 diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index f126d477cba6..b0ec0f1830ca 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -980,9 +980,9 @@ bool DepthToSpaceRel(const Array& types, int num_inputs, const Attrs& attr << " But got " << in_layout; auto oshape = layout_converter.ForwardShape(data->shape); - oshape.Set(1, indexdiv(data->shape[1], (block_size * block_size))); - oshape.Set(2, data->shape[2] * block_size); - oshape.Set(3, data->shape[3] * block_size); + 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], From 4d72ebadba6fd2be7aa7b2642061af2dc5c65c00 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Sun, 22 Dec 2019 10:22:08 -0800 Subject: [PATCH 07/11] SpaceToDepth in and all tests passing. --- python/tvm/relay/frontend/onnx.py | 51 +--------------- python/tvm/relay/frontend/tensorflow.py | 33 +---------- python/tvm/relay/op/nn/_nn.py | 11 ++++ python/tvm/relay/op/nn/nn.py | 23 ++++++++ src/relay/op/nn/nn.cc | 59 ++++++++++++++++++- .../frontend/tensorflow/test_forward.py | 1 - tests/python/relay/test_op_level5.py | 32 ++++++++++ 7 files changed, 128 insertions(+), 82 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 52211f89221a..dbe2b2af0f74 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -539,35 +539,7 @@ class DepthToSpace(OnnxOpConverter): 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) class SpaceToDepth(OnnxOpConverter): @@ -578,26 +550,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 27a4d3c8d33e..787d5cbd59bc 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -673,38 +673,9 @@ def _impl(inputs, attr, params): 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 5b20bed82b19..148b9d2f78da 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -891,6 +891,17 @@ def compute_depth_to_space(attrs, inputs, out_dtype, target): 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 06354687afc8..20a7b9ca3142 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -2008,3 +2008,26 @@ def depth_to_space(data, block_size, layout='NCHW'): in_height * block_size, in_width * block_size] """ return _make.depth_to_space(data, block_size, layout) + + +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) \ No newline at end of file diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index b0ec0f1830ca..4489bf4c9632 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -991,7 +991,7 @@ bool DepthToSpaceRel(const Array& types, int num_inputs, const Attrs& attr return true; } -// Positional relay function to create DepthToSpace opeator +// Positional relay function to create DepthToSpace operator // used by frontend FFI Expr MakeDepthToSpace(Expr data, int block_size, std::string layout) { auto attrs = make_node(); @@ -1019,5 +1019,62 @@ RELAY_REGISTER_OP("nn.depth_to_space") .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/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 82de233f7b7e..090573f32ea7 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -2855,7 +2855,6 @@ def test_forward_add_n(): test_forward_sin() test_forward_negative() test_forward_divide() - test_forward_floordiv() test_forward_abs() test_forward_softplus() test_forward_sqrt() diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index 3910a20be092..da23d1c31d9a 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -604,6 +604,37 @@ def verify_depth_to_space(dshape, block_size, layout): verify_depth_to_space((1, 4, 4, 4), 2, layout) +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() @@ -618,3 +649,4 @@ def verify_depth_to_space(dshape, block_size, layout): test_non_max_suppression() test_deformable_conv2d() test_depth_to_space() + test_space_to_depth() \ No newline at end of file From 2704f8cb77a62c9c3b4b1bd759613fa2a2620074 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Sun, 22 Dec 2019 10:26:55 -0800 Subject: [PATCH 08/11] lint fixes. --- python/tvm/relay/op/nn/nn.py | 10 +++++----- python/tvm/relay/op/op_attrs.py | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index 20a7b9ca3142..ed7a424feda0 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -1994,7 +1994,7 @@ def depth_to_space(data, block_size, layout='NCHW'): ---------- data : tvm.relay.Expr Input data with channels divisible by block_size**2 - + block_size : int Size of blocks to convert channels into. @@ -2004,7 +2004,7 @@ def depth_to_space(data, block_size, layout='NCHW'): Returns ------- result : tvm.relay.Expr - Tensor with shape [in_batch, in_channel / block_size * block_size, + 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) @@ -2017,7 +2017,7 @@ def space_to_depth(data, block_size, layout='NCHW'): ---------- data : tvm.relay.Expr Input data with spatial dimensions divisible by block_size - + block_size : int Size of blocks to decompose into channels. @@ -2027,7 +2027,7 @@ def space_to_depth(data, block_size, layout='NCHW'): Returns ------- result : tvm.relay.Expr - Tensor with shape [in_batch, in_channel * block_size * block_size, + 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) \ No newline at end of file + 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 44ce60315bda..f8cafd30166d 100644 --- a/python/tvm/relay/op/op_attrs.py +++ b/python/tvm/relay/op/op_attrs.py @@ -293,4 +293,4 @@ class Conv2DTransposeAttrs(Attrs): @register_relay_attr_node class SubPixelAttrs(Attrs): - """Attributes used in depth to space and space to depth operators""" \ No newline at end of file + """Attributes used in depth to space and space to depth operators""" From 4918e8e7e9e637faebefc27a11cfed583afea70e Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Sun, 22 Dec 2019 14:02:41 -0500 Subject: [PATCH 09/11] Added string include --- src/relay/op/nn/nn.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index 4489bf4c9632..2a9b16653c0a 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" From 0dd39719602713b66131b952f2968c2f95496c18 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Sun, 22 Dec 2019 14:22:30 -0500 Subject: [PATCH 10/11] Fixed topi formatting. --- topi/python/topi/nn/depth_to_space.py | 40 +++++++++++++---- topi/python/topi/nn/space_to_depth.py | 45 ++++++++++++++----- topi/python/topi/testing/depth_to_space.py | 2 +- topi/python/topi/testing/space_to_depth.py | 2 +- topi/tests/python/test_topi_depth_to_space.py | 3 +- topi/tests/python/test_topi_space_to_depth.py | 3 +- 6 files changed, 72 insertions(+), 23 deletions(-) diff --git a/topi/python/topi/nn/depth_to_space.py b/topi/python/topi/nn/depth_to_space.py index bedd806204fc..e55b9e4189d3 100644 --- a/topi/python/topi/nn/depth_to_space.py +++ b/topi/python/topi/nn/depth_to_space.py @@ -20,39 +20,61 @@ import tvm from .. import tag + def depth_to_space(data, block_size, layout='NCHW'): + """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. + + 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] + 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] + 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) channel_idx = channel_factor * ((block_size * idx_y) + idx_x) + c - + if layout == 'NCHW': - return data(n, channel_idx, block_y, block_x) + output = data(n, channel_idx, block_y, block_x) else: - return data(n, block_y, block_x, channel_idx) - + 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 index d8d8a23b40e8..6ed7cd64a448 100644 --- a/topi/python/topi/nn/space_to_depth.py +++ b/topi/python/topi/nn/space_to_depth.py @@ -20,36 +20,61 @@ 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)] + 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] + 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 + 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': - return data(n, channel_idx, y_idx + (y * block_size), x_idx + (x * block_size)) + output = data(n, channel_idx, y_idx + + (y * block_size), x_idx + (x * block_size)) else: - return data(n, y_idx + (y * block_size), x_idx + (x * block_size), channel_idx) - + 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/depth_to_space.py b/topi/python/topi/testing/depth_to_space.py index c3eab104c225..d361f72ae5f0 100644 --- a/topi/python/topi/testing/depth_to_space.py +++ b/topi/python/topi/testing/depth_to_space.py @@ -33,7 +33,7 @@ def depth_to_space_python(data, block_size): Returns ------- d2s_out : np.ndarray - 4-D with shape [batch, in_channel / (block_size * block_size), + 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 diff --git a/topi/python/topi/testing/space_to_depth.py b/topi/python/topi/testing/space_to_depth.py index 5a1a2dc9a4fe..3a3b94177602 100644 --- a/topi/python/topi/testing/space_to_depth.py +++ b/topi/python/topi/testing/space_to_depth.py @@ -33,7 +33,7 @@ def space_to_depth_python(data, block_size): Returns ------- d2s_out : np.ndarray - 4-D with shape [batch, in_channel * (block_size * block_size), + 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 diff --git a/topi/tests/python/test_topi_depth_to_space.py b/topi/tests/python/test_topi_depth_to_space.py index c4694dd27293..ed5e65e0108c 100644 --- a/topi/tests/python/test_topi_depth_to_space.py +++ b/topi/tests/python/test_topi_depth_to_space.py @@ -66,6 +66,7 @@ def check_device(device): for device in get_all_backend(): check_device(device) + def test_depth_to_space(): for layout in ['NCHW', 'NHWC']: # Simplest possible case @@ -81,4 +82,4 @@ def test_depth_to_space(): if __name__ == "__main__": - test_depth_to_space() \ No newline at end of file + 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 index 12e93133e495..b25cad194301 100644 --- a/topi/tests/python/test_topi_space_to_depth.py +++ b/topi/tests/python/test_topi_space_to_depth.py @@ -66,6 +66,7 @@ def check_device(device): for device in get_all_backend(): check_device(device) + def test_space_to_depth(): for layout in ['NCHW', 'NHWC']: # Simplest possible case @@ -81,4 +82,4 @@ def test_space_to_depth(): if __name__ == "__main__": - test_space_to_depth() \ No newline at end of file + test_space_to_depth() From f4f10b1aad6825381879bdc301f883a1ab3fbc31 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Sun, 22 Dec 2019 23:28:29 -0500 Subject: [PATCH 11/11] Added DCR/CDR mode to depthtospace operator. --- include/tvm/relay/attrs/nn.h | 14 +++++----- python/tvm/relay/frontend/onnx.py | 3 ++- python/tvm/relay/op/nn/_nn.py | 3 ++- python/tvm/relay/op/nn/nn.py | 8 ++++-- src/relay/op/nn/nn.cc | 3 ++- tests/python/relay/test_op_level5.py | 9 ++++--- topi/python/topi/nn/depth_to_space.py | 12 +++++++-- topi/python/topi/testing/depth_to_space.py | 13 ++++++--- topi/tests/python/test_topi_depth_to_space.py | 27 ++++++++++--------- 9 files changed, 58 insertions(+), 34 deletions(-) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index 354349624f58..4e061bd7d256 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -825,17 +825,19 @@ struct DeformableConv2DAttrs : public tvm::AttrsNode { 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(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 diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index dbe2b2af0f74..254b34bca4a7 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -539,7 +539,8 @@ class DepthToSpace(OnnxOpConverter): def _impl_v11(cls, inputs, attr, params): block_size = int(attr['blocksize']) - return _op.nn.depth_to_space(inputs[0], block_size) + mode = attr.get("mode", "DCR") + return _op.nn.depth_to_space(inputs[0], block_size, mode=mode) class SpaceToDepth(OnnxOpConverter): diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index 148b9d2f78da..c97eede8a993 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -886,7 +886,8 @@ def compute_cross_entropy_with_logits(attrs, inputs, out_dtype, target): def compute_depth_to_space(attrs, inputs, out_dtype, target): block_size = attrs.block_size layout = attrs.layout - return [topi.nn.depth_to_space(inputs[0], block_size, layout=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) diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index ed7a424feda0..95096af1d32f 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -1987,7 +1987,7 @@ def cross_entropy_with_logits(predictions, targets): return _make.cross_entropy_with_logits(predictions, targets) -def depth_to_space(data, block_size, layout='NCHW'): +def depth_to_space(data, block_size, layout='NCHW', mode='DCR'): """Convert channels into spatial blocks. Parameters @@ -2001,13 +2001,17 @@ def depth_to_space(data, block_size, layout='NCHW'): 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) + return _make.depth_to_space(data, block_size, layout, mode) def space_to_depth(data, block_size, layout='NCHW'): diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index 2a9b16653c0a..ac3848513856 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -994,10 +994,11 @@ bool DepthToSpaceRel(const Array& types, int num_inputs, const Attrs& attr // Positional relay function to create DepthToSpace operator // used by frontend FFI -Expr MakeDepthToSpace(Expr data, int block_size, std::string layout) { +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), {}); } diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index da23d1c31d9a..84e9f55d67e7 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -574,7 +574,7 @@ def test_run(batch, in_channel, size, out_channel, deformable_groups, groups): def test_depth_to_space(): - def verify_depth_to_space(dshape, block_size, layout): + 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: @@ -583,13 +583,13 @@ def verify_depth_to_space(dshape, block_size, layout): 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) + 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) + 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") @@ -601,7 +601,8 @@ def verify_depth_to_space(dshape, block_size, layout): 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_depth_to_space((1, 4, 4, 4), 2, layout) + for mode in ["DCR", "CDR"]: + verify_depth_to_space((1, 4, 4, 4), 2, layout, mode) def test_space_to_depth(): diff --git a/topi/python/topi/nn/depth_to_space.py b/topi/python/topi/nn/depth_to_space.py index e55b9e4189d3..d847c08daf27 100644 --- a/topi/python/topi/nn/depth_to_space.py +++ b/topi/python/topi/nn/depth_to_space.py @@ -21,7 +21,7 @@ from .. import tag -def depth_to_space(data, block_size, layout='NCHW'): +def depth_to_space(data, block_size, layout='NCHW', mode='DCR'): """Perform depth to space transformation on the data Parameters @@ -35,6 +35,11 @@ def depth_to_space(data, block_size, layout='NCHW'): 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 @@ -65,7 +70,10 @@ def _get_pixel(n, c, y, x): block_y = tvm.truncdiv(y, block_size) idx_x = tvm.truncmod(x, block_size) idx_y = tvm.truncmod(y, block_size) - channel_idx = channel_factor * ((block_size * idx_y) + idx_x) + c + 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) diff --git a/topi/python/topi/testing/depth_to_space.py b/topi/python/topi/testing/depth_to_space.py index d361f72ae5f0..f4a60bc12d87 100644 --- a/topi/python/topi/testing/depth_to_space.py +++ b/topi/python/topi/testing/depth_to_space.py @@ -19,7 +19,7 @@ import numpy as np -def depth_to_space_python(data, block_size): +def depth_to_space_python(data, block_size, mode='DCR'): """Depth to Space operator in python for NCHW layout. Parameters @@ -41,9 +41,14 @@ def depth_to_space_python(data, 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, block_size, block_size, new_c, in_h, in_w]) - transposed = np.transpose(expanded, axes=[0, 3, 4, 1, 5, 2]) + 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/tests/python/test_topi_depth_to_space.py b/topi/tests/python/test_topi_depth_to_space.py index ed5e65e0108c..4e895cb5db55 100644 --- a/topi/tests/python/test_topi_depth_to_space.py +++ b/topi/tests/python/test_topi_depth_to_space.py @@ -23,7 +23,7 @@ from common import get_all_backend -def verify_depth_to_space(block_size, batch, in_channel, in_height, in_width, layout='NCHW'): +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) @@ -41,10 +41,10 @@ def verify_depth_to_space(block_size, batch, in_channel, in_height, in_width, la 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) + 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) + 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]) @@ -69,16 +69,17 @@ def check_device(device): def test_depth_to_space(): for layout in ['NCHW', 'NHWC']: - # Simplest possible case - verify_depth_to_space(2, 1, 4, 1, 1, layout=layout) - # Average input size - verify_depth_to_space(2, 1, 32, 32, 32, layout=layout) - # Large block size - verify_depth_to_space(8, 1, 256, 32, 32, layout=layout) - # Large batch size - verify_depth_to_space(4, 8, 32, 32, 32, layout=layout) - # Large input size - verify_depth_to_space(4, 8, 32, 128, 128, layout=layout) + 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__":