diff --git a/docs/api/python/topi.rst b/docs/api/python/topi.rst index b1f2bfc9..3216c516 100644 --- a/docs/api/python/topi.rst +++ b/docs/api/python/topi.rst @@ -51,6 +51,7 @@ List of operators topi.broadcast_div topi.broadcast_maximum topi.broadcast_minimum + topi.image.resize List of schedules @@ -114,6 +115,10 @@ topi.nn .. autofunction:: topi.nn.depthwise_conv2d_nchw .. autofunction:: topi.nn.depthwise_conv2d_nhwc +topi.image +~~~~~~~~~~ +.. autofunction:: topi.image.resize + topi.generic ~~~~~~~~~~~~ diff --git a/nnvm/include/nnvm/top/nn.h b/nnvm/include/nnvm/top/nn.h index 7593cf88..bee61378 100644 --- a/nnvm/include/nnvm/top/nn.h +++ b/nnvm/include/nnvm/top/nn.h @@ -288,16 +288,22 @@ struct GlobalPool2DParam : public dmlc::Parameter { struct UpSamplingParam : public dmlc::Parameter { int scale; std::string layout; + std::string method; DMLC_DECLARE_PARAMETER(UpSamplingParam) { DMLC_DECLARE_FIELD(scale) .describe("upsampling scaling factor"); DMLC_DECLARE_FIELD(layout) .set_default("NCHW") - .describe("Dimension ordering of data and weight. Can be 'NCHW', 'NHWC', etc." + .describe("Dimension ordering of data. Can be 'NCHW', 'NHWC', etc." "'N', 'C', 'H', 'W' stands for batch, channel, height, and width" - "dimensions respectively. Convolution is applied on the 'H' and" + "dimensions respectively. Upsampling is applied on the 'H' and" "'W' dimensions."); + DMLC_DECLARE_FIELD(method) + .set_default("NEAREST_NEIGHBOR") + .describe("Specify the mode to use for scaling." + "NEAREST_NEIGHBOR - Nearest Neighbor" + "BILINEAR - Bilinear Interpolation"); } }; diff --git a/nnvm/python/nnvm/top/__init__.py b/nnvm/python/nnvm/top/__init__.py index 12294fa0..d9174f81 100644 --- a/nnvm/python/nnvm/top/__init__.py +++ b/nnvm/python/nnvm/top/__init__.py @@ -8,6 +8,7 @@ from . import nn from . import transform from . import reduction from . import vision +from . import image from .registry import OpPattern from .registry import register_compute, register_schedule, register_pattern diff --git a/nnvm/python/nnvm/top/image.py b/nnvm/python/nnvm/top/image.py new file mode 100644 index 00000000..f903231e --- /dev/null +++ b/nnvm/python/nnvm/top/image.py @@ -0,0 +1,17 @@ +# pylint: disable=invalid-name, unused-argument +"""Definition of image ops""" +from __future__ import absolute_import + +import topi +import tvm +from . import registry as reg +from .registry import OpPattern + +# resize +@reg.register_schedule("resize") +def schedule_resize(_, outs, target): + """Schedule definition of resize""" + with tvm.target.create(target): + return topi.generic.schedule_injective(outs) + +reg.register_pattern("resize", OpPattern.INJECTIVE) diff --git a/nnvm/python/nnvm/top/nn.py b/nnvm/python/nnvm/top/nn.py index dfebc368..b7e0d095 100644 --- a/nnvm/python/nnvm/top/nn.py +++ b/nnvm/python/nnvm/top/nn.py @@ -235,20 +235,10 @@ def schedule_global_avg_pool2d(_, outs, target): reg.register_pattern("global_avg_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE) - -@reg.register_compute("upsampling") -def compute_upsampling(attrs, inputs, _): - """Compute definition of upsampling""" - scale = attrs.get_int("scale") - layout = attrs["layout"] - if layout: - assert layout == "NCHW" or layout == "NHWC" - return topi.nn.upsampling(inputs[0], scale, layout) - return topi.nn.upsampling(inputs[0], scale) - +# upsampling @reg.register_schedule("upsampling") def schedule_upsampling(_, outs, target): - """Compute definition of upsampling""" + """Schedule definition of upsampling""" with tvm.target.create(target): return topi.generic.schedule_injective(outs) diff --git a/nnvm/src/top/image/resize.cc b/nnvm/src/top/image/resize.cc new file mode 100644 index 00000000..151296d8 --- /dev/null +++ b/nnvm/src/top/image/resize.cc @@ -0,0 +1,113 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file resize.cc + * \brief Property def of resize operators. + */ +#include +#include +#include +#include +#include +#include +#include +#include +#include "../nn/nn_common.h" +#include "../op_common.h" +#include "../elemwise_op_common.h" +#include "topi/elemwise.h" +#include "topi/transform.h" +#include "topi/image/resize.h" +#include "resize.h" + +namespace nnvm { +namespace top { +using tvm::Expr; +using tvm::Array; +using tvm::Tensor; +using nnvm::compiler::FTVMCompute; + +DMLC_REGISTER_PARAMETER(ResizeParam); + +inline bool ResizeInferShape(const nnvm::NodeAttrs& attrs, + std::vector* in_shape, + std::vector* out_shape) { + static const Layout kNCHW("NCHW"); + const ResizeParam& param = nnvm::get(attrs.parsed); + CHECK_EQ(in_shape->size(), 1U); + CHECK_EQ(out_shape->size(), 1U); + TShape dshape = (*in_shape)[0]; + if (dshape.ndim() == 0) return false; + dshape = ConvertLayout(dshape, param.layout, kNCHW); + + TShape oshape = dshape; + if (param.layout == "NCHW") { + oshape[2] = param.size[0]; + oshape[3] = param.size[1]; + } else { + oshape[1] = param.size[0]; + oshape[2] = param.size[1]; + } + oshape = ConvertLayout(oshape, kNCHW, param.layout); + NNVM_ASSIGN_OUTPUT_SHAPE(attrs, *out_shape, 0, oshape); + + return true; +} + +inline bool ResizeLayout(const NodeAttrs& attrs, + std::vector *in_layouts, + const std::vector *last_in_layouts, + std::vector *out_layouts) { + const ResizeParam& param = nnvm::get(attrs.parsed); + CHECK_EQ(in_layouts->size(), 1U); + CHECK_EQ(out_layouts->size(), 1U); + const Layout layout(param.layout); + NNVM_ASSIGN_LAYOUT(*in_layouts, 0, layout); + NNVM_ASSIGN_LAYOUT(*out_layouts, 0, layout); + return true; +} + +NNVM_REGISTER_OP(resize) +.describe(R"(Perform resize to input array with nearest neighbour or bilinear interpolation. + +- **data**: data is 4D array of shape + (batch_size, channels, in_height, in_width) for NCHW + (batch_size, in_height, in_width, channels) for NHWC + +- **out**: Output is 4D array of shape + for layout NCHW + (batch_size, channels, size[0], size[1]) + + for layout NHWC + (batch_size, size[0], size[1], channels) + +)" NNVM_ADD_FILELINE) +.add_argument("data", "4D Tensor", "Input data.") +.add_arguments(ResizeParam::__FIELDS__()) +.set_attr_parser(ParamParser) +.set_attr("FGetAttrDict", ParamGetAttrDict) +.set_attr("FInferShape", ResizeInferShape) +.set_attr("FInferType", ElemwiseType<1, 1>) +.set_attr("FCorrectLayout", ResizeLayout) +.set_num_outputs(1) +.set_num_inputs(1) +.set_attr( + "FTVMCompute", [](const NodeAttrs& attrs, + const Array& inputs, + const Array& out_info) { + const ResizeParam& param = nnvm::get(attrs.parsed); + Array oshape; + if (param.layout == "NCHW") { + oshape.push_back(out_info[0]->shape[2]); + oshape.push_back(out_info[0]->shape[3]); + } else { + oshape.push_back(out_info[0]->shape[1]); + oshape.push_back(out_info[0]->shape[2]); + } + + return Array{ topi::image::resize(inputs[0], oshape, param.layout, + param.align_corners, param.method)}; +}) +.set_support_level(2); + +} // namespace top +} // namespace nnvm diff --git a/nnvm/src/top/image/resize.h b/nnvm/src/top/image/resize.h new file mode 100644 index 00000000..7237152f --- /dev/null +++ b/nnvm/src/top/image/resize.h @@ -0,0 +1,45 @@ +/*! + * Copyright (c) 2018 by Contributors + * \file resize.h + */ +#ifndef NNVM_TOP_IMAGE_RESIZE_H_ +#define NNVM_TOP_IMAGE_RESIZE_H_ + +#include +#include +#include +#include +#include + +namespace nnvm { +namespace top { + +struct ResizeParam : public dmlc::Parameter { + TShape size; + std::string layout; + std::string method; + bool align_corners; + + DMLC_DECLARE_PARAMETER(ResizeParam) { + DMLC_DECLARE_FIELD(size) + .describe("Output size"); + DMLC_DECLARE_FIELD(layout) + .set_default("NCHW") + .describe("Dimension ordering of data. Can be 'NCHW', 'NHWC', etc." + "'N', 'C', 'H', 'W' stands for batch, channel, height, and width" + "dimensions respectively. Resize is applied on the 'H' and" + "'W' dimensions."); + DMLC_DECLARE_FIELD(method) + .set_default("BILINEAR") + .describe("Specify the mode to use for scaling." + "NEAREST_NEIGHBOR - Nearest Neighbor" + "BILINEAR - Bilinear Interpolation"); + DMLC_DECLARE_FIELD(align_corners) + .set_default(false) + .describe("Should be true to preserve the values at the corner pixels"); + } +}; + +} // namespace top +} // namespace nnvm +#endif // NNVM_TOP_IMAGE_RESIZE_H_ diff --git a/nnvm/src/top/nn/upsampling.cc b/nnvm/src/top/nn/upsampling.cc index 73c9bad9..6c5e1344 100644 --- a/nnvm/src/top/nn/upsampling.cc +++ b/nnvm/src/top/nn/upsampling.cc @@ -1,8 +1,12 @@ /*! * Copyright (c) 2017 by Contributors - * \file pooling.cc - * \brief Property def of pooling operators. + * \file upsampling.cc + * \brief Property def of upsampling operators. */ +#include +#include +#include +#include #include #include #include @@ -10,27 +14,36 @@ #include "./nn_common.h" #include "../op_common.h" #include "../elemwise_op_common.h" +#include "topi/elemwise.h" +#include "topi/transform.h" +#include "topi/nn/upsampling.h" namespace nnvm { namespace top { +using tvm::Expr; +using tvm::Array; +using tvm::Tensor; +using nnvm::compiler::FTVMCompute; DMLC_REGISTER_PARAMETER(UpSamplingParam); inline bool UpSamplingInferShape(const nnvm::NodeAttrs& attrs, - std::vector* in_shape, - std::vector* out_shape) { + std::vector* in_shape, + std::vector* out_shape) { static const Layout kNCHW("NCHW"); const UpSamplingParam& param = nnvm::get(attrs.parsed); CHECK_EQ(in_shape->size(), 1U); CHECK_EQ(out_shape->size(), 1U); TShape dshape = (*in_shape)[0]; if (dshape.ndim() == 0) return false; + dshape = ConvertLayout(dshape, param.layout, kNCHW); TShape oshape = dshape; oshape[2] = oshape[2] * param.scale; oshape[3] = oshape[3] * param.scale; oshape = ConvertLayout(oshape, kNCHW, param.layout); NNVM_ASSIGN_OUTPUT_SHAPE(attrs, *out_shape, 0, oshape); + return true; } @@ -48,10 +61,18 @@ inline bool UpsamplingLayout(const NodeAttrs& attrs, } NNVM_REGISTER_OP(upsampling) -.describe(R"(Perform nearest neighbor upsampling to input array. +.describe(R"(Perform upsampling to input array with nearest neighbour or bilinear interpolation. -- **data**: Input is 4D array of shape (batch_size, channels, in_height, in_width). -- **out**: Output is 4D array of shape (batch_size, channels, in_height*scale, in_width*scale). +- **data**: data is 4D array of shape + (batch_size, channels, in_height, in_width) for NCHW + (batch_size, in_height, in_width, channels) for NHWC + +- **out**: Output is 4D array of shape + for layout NCHW + (batch_size, channels, in_height*scale, in_width*scale) + + for layout NHWC + (batch_size, in_height*scale, in_width*scale, channels) )" NNVM_ADD_FILELINE) .add_argument("data", "4D Tensor", "Input data.") @@ -63,6 +84,22 @@ NNVM_REGISTER_OP(upsampling) .set_attr("FCorrectLayout", UpsamplingLayout) .set_num_outputs(1) .set_num_inputs(1) +.set_attr( + "FTVMCompute", [](const NodeAttrs& attrs, + const Array& inputs, + const Array& out_info) { + const UpSamplingParam& param = nnvm::get(attrs.parsed); + Array oshape; + if (param.layout == "NCHW") { + oshape.push_back(out_info[0]->shape[2]); + oshape.push_back(out_info[0]->shape[3]); + } else { + oshape.push_back(out_info[0]->shape[1]); + oshape.push_back(out_info[0]->shape[2]); + } + + return Array{ topi::nn::upsampling(inputs[0], oshape, param.layout, param.method)}; +}) .set_support_level(2); } // namespace top diff --git a/nnvm/tests/python/compiler/test_top_level2.py b/nnvm/tests/python/compiler/test_top_level2.py index f752b7d4..ed67dab4 100644 --- a/nnvm/tests/python/compiler/test_top_level2.py +++ b/nnvm/tests/python/compiler/test_top_level2.py @@ -210,7 +210,7 @@ def test_global_avg_pool2d(): np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5) -def test_upsampling(): +def test_upsampling_nearest_neighbor(): x = sym.Variable("x") scale = 2 y = sym.upsampling(x, scale=scale, name="y") @@ -225,9 +225,46 @@ def test_upsampling(): data = tvm.nd.array(a_np) m.run(x=data) out = m.get_output(0, tvm.nd.empty(oshape, dtype)) - b_np = topi.testing.upsampling_python(a_np, scale) + b_np = topi.testing.upsampling_python(a_np, scale, "NCHW") np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5) +def test_upsampling_bilinear(): + x = sym.Variable("x") + scale = 2 + y = sym.upsampling(x, scale=scale, method="BILINEAR", name="y", layout="NCHW") + dtype = "float32" + dshape = (1, 4, 32, 32) + oshape = (1, 4, 32*scale, 32*scale) + shape_dict = {"x": dshape} + dtype_dict = {"x": dtype} + for target, ctx in ctx_list(): + graph, lib, _ = nnvm.compiler.build(y, target, shape_dict, dtype_dict) + m = graph_runtime.create(graph, lib, ctx) + a_np = np.random.uniform(size=dshape).astype(dtype) + data = tvm.nd.array(a_np) + m.run(x=data) + out = m.get_output(0, tvm.nd.empty(oshape, dtype)) + b_np = topi.testing.bilinear_resize_python(a_np, (32*scale, 32*scale), "NCHW") + np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5, atol=1e-5) + +def test_resize_bilinear(): + x = sym.Variable("x") + scale = 2 + y = sym.upsampling(x, scale=scale, method="BILINEAR", name="y", layout="NHWC") + dtype = "float32" + dshape = (1, 32, 32, 4) + oshape = (1, 32*scale, 32*scale, 4) + shape_dict = {"x": dshape} + dtype_dict = {"x": dtype} + for target, ctx in ctx_list(): + graph, lib, _ = nnvm.compiler.build(y, target, shape_dict, dtype_dict) + m = graph_runtime.create(graph, lib, ctx) + a_np = np.random.uniform(size=dshape).astype(dtype) + data = tvm.nd.array(a_np) + m.run(x=data) + out = m.get_output(0, tvm.nd.empty(oshape, dtype)) + b_np = topi.testing.bilinear_resize_python(a_np, (32*scale, 32*scale), "NHWC") + np.testing.assert_allclose(out.asnumpy(), b_np, rtol=1e-5, atol=1e-5) if __name__ == "__main__": test_conv2d() @@ -239,4 +276,6 @@ if __name__ == "__main__": test_avg_pool2d_no_count_pad() test_global_max_pool2d() test_global_avg_pool2d() - test_upsampling() + test_upsampling_nearest_neighbor() + test_upsampling_bilinear() + test_resize_bilinear() diff --git a/topi/include/topi/image/resize.h b/topi/include/topi/image/resize.h new file mode 100644 index 00000000..e00fe15f --- /dev/null +++ b/topi/include/topi/image/resize.h @@ -0,0 +1,316 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file topi/image/resize.h + * \brief image resize constructors + */ +#ifndef TOPI_IMAGE_RESIZE_H_ +#define TOPI_IMAGE_RESIZE_H_ + +#include +#include +#include +#include + +#include "topi/tags.h" +#include "topi/detail/ravel_unravel.h" +#include "topi/detail/constant_utils.h" +#include "tvm/tvm.h" + +namespace topi { +namespace image { +using namespace tvm; + +/*! +* \brief Resize given tensor to given shape using nearest neighbour for NHWC +* +* \param input The input tensor. +* \param shape Output shape to resize to. +* \param align_corners To preserve centers of 4 corner pixels +* \param name Name of the operation +* \param tag The tag to mark the operation +* +* \return A Tensor resized to given shape +*/ +inline Tensor resize_nearest_neighbor_nhwc(const Tensor& input, + const Array& shape, + bool align_corners = false, + std::string name = "tensor", + std::string tag = kInjective) { + Array out_shape; + out_shape.push_back(input->shape[0]); + out_shape.push_back(shape[0]); + out_shape.push_back(shape[1]); + out_shape.push_back(input->shape[3]); + + Expr h_ratio = shape[0] / input->shape[1]; + Expr w_ratio = shape[1] / input->shape[2]; + + return compute( + out_shape, [&](const Array& indices) { + Array idx; + idx.push_back(indices[0]); + idx.push_back(indices[1] / h_ratio); + idx.push_back(indices[2] / w_ratio); + idx.push_back(indices[3]); + + return input(idx); + }, name, tag); +} + +/*! +* \brief Resize given tensor to given shape using nearest neighbour for NCHW +* +* \param input The input tensor. +* \param shape Output shape to resize to. +* \param align_corners To preserve centers of 4 corner pixels +* \param name Name of the operation +* \param tag The tag to mark the operation +* +* \return A Tensor resized to given shape +*/ +inline Tensor resize_nearest_neighbor_nchw(const Tensor& input, + const Array& shape, + bool align_corners = false, + std::string name = "tensor", + std::string tag = kInjective) { + Array out_shape; + out_shape.push_back(input->shape[0]); + out_shape.push_back(input->shape[1]); + out_shape.push_back(shape[0]); + out_shape.push_back(shape[1]); + + Expr h_ratio = shape[0] / input->shape[2]; + Expr w_ratio = shape[1] / input->shape[3]; + + return compute( + out_shape, [&](const Array& indices) { + Array idx; + idx.push_back(indices[0]); + idx.push_back(indices[1]); + idx.push_back(indices[2] / h_ratio); + idx.push_back(indices[3] / w_ratio); + + return input(idx); + }, name, tag); +} + +/*! +* \brief Resize given tensor to given shape using nearest neighbour +* +* \param input The input tensor. +* \param shape Output shape to resize to. +* \param layout input layout +* \param align_corners To preserve centers of 4 corner pixels +* \param name Name of the operation +* \param tag The tag to mark the operation +* +* \return A Tensor resized to given shape +*/ +inline Tensor resize_nearest_neighbor(const Tensor& input, + const Array& shape, + std::string layout = "NCHW", + bool align_corners = false, + std::string name = "tensor", + std::string tag = kInjective) { + CHECK_EQ(align_corners, false) << "Align corners not supported for nearest neighbour"; + + if (layout == "NHWC") { + return resize_nearest_neighbor_nhwc(input, shape, align_corners); + } else { + return resize_nearest_neighbor_nchw(input, shape, align_corners); + } +} + +/*! +* \brief Resize given tensor to given shape using bilinear interpolation for NHWC +* +* \param input The input tensor. +* \param shape Output shape to resize to. +* \param align_corners To preserve centers of 4 corner pixels +* \param name Name of the operation +* \param tag The tag to mark the operation +* +* \return A Tensor resized to given shape +*/ +inline Tensor resize_bilinear_nhwc(const Tensor& input, + const Array& shape, + bool align_corners = false, + std::string name = "tensor", + std::string tag = kInjective) { + Array out_shape; + out_shape.push_back(input->shape[0]); + out_shape.push_back(shape[0]); + out_shape.push_back(shape[1]); + out_shape.push_back(input->shape[3]); + + Expr cone = make_const(Int(32), 1); + + auto in_height = as_const_int(input->shape[1]); + auto in_width = as_const_int(input->shape[2]); + auto out_height = as_const_int(shape[0]); + auto out_width = as_const_int(shape[1]); + + Expr y_ratio; + Expr x_ratio; + + if (align_corners) { + y_ratio = make_const(Float(32), (static_cast(*in_height) / + static_cast(*out_height))); + x_ratio = make_const(Float(32), (static_cast(*in_width) / + static_cast(*out_width))); + } else { + y_ratio = make_const(Float(32), (static_cast(*in_height - 1) / + static_cast(*out_height - 1))); + x_ratio = make_const(Float(32), (static_cast(*in_width - 1) / + static_cast(*out_width - 1))); + } + + Expr other_y = tvm::ir::Simplify(input->shape[1] - cone); + Expr other_x = tvm::ir::Simplify(input->shape[2] - cone); + + return compute( + out_shape, [&](const Array& indices) { + auto y0 = HalideIR::Internal::Cast::make(Int(32), tvm::floor(y_ratio * indices[1])); + auto x0 = HalideIR::Internal::Cast::make(Int(32), tvm::floor(x_ratio * indices[2])); + + auto y1 = tvm::select(((y0 + cone) > other_y), other_y, (y0 + cone)); + auto x1 = tvm::select(((x0 + cone) > other_x), other_x, (x0 + cone)); + + auto h = (y_ratio * indices[1]) - y0; + auto w = (x_ratio * indices[2]) - x0;; + + auto A = input(indices[0], y0, x0, indices[3]); + auto B = input(indices[0], y0, x1, indices[3]); + auto C = input(indices[0], y1, x0, indices[3]); + auto D = input(indices[0], y1, x1, indices[3]); + + return (A*(cone-w)*(cone-h) + B*(w)*(cone-h) + C*(h)*(cone-w) + D*w*h); + }, name, tag); +} + +/*! +* \brief Resize given tensor to given shape using bilinear interpolation for NCHW +* +* \param input The input tensor. +* \param shape Output shape to resize to. +* \param align_corners To preserve centers of 4 corner pixels +* \param name Name of the operation +* \param tag The tag to mark the operation +* +* \return A Tensor resized to given shape +*/ +inline Tensor resize_bilinear_nchw(const Tensor& input, + const Array& shape, + bool align_corners = false, + std::string name = "tensor", + std::string tag = kInjective) { + Array out_shape; + out_shape.push_back(input->shape[0]); + out_shape.push_back(input->shape[1]); + out_shape.push_back(shape[0]); + out_shape.push_back(shape[1]); + + Expr cone = make_const(Int(32), 1); + + auto in_height = as_const_int(input->shape[2]); + auto in_width = as_const_int(input->shape[3]); + auto out_height = as_const_int(shape[0]); + auto out_width = as_const_int(shape[1]); + + Expr y_ratio; + Expr x_ratio; + + if (align_corners) { + y_ratio = make_const(Float(32), (static_cast(*in_height) / + static_cast(*out_height))); + x_ratio = make_const(Float(32), (static_cast(*in_width) / + static_cast(*out_width))); + } else { + y_ratio = make_const(Float(32), (static_cast(*in_height - 1) / + static_cast(*out_height - 1))); + x_ratio = make_const(Float(32), (static_cast(*in_width - 1) / + static_cast(*out_width - 1))); + } + + Expr other_y = tvm::ir::Simplify(input->shape[2] - cone); + Expr other_x = tvm::ir::Simplify(input->shape[3] - cone); + + return compute( + out_shape, [&](const Array& indices) { + auto y0 = HalideIR::Internal::Cast::make(Int(32), tvm::floor(y_ratio * indices[2])); + auto x0 = HalideIR::Internal::Cast::make(Int(32), tvm::floor(x_ratio * indices[3])); + + auto y1 = tvm::select(((y0 + cone) > other_y), other_y, (y0 + cone)); + auto x1 = tvm::select(((x0 + cone) > other_x), other_x, (x0 + cone)); + + auto h = (y_ratio * indices[2]) - y0; + auto w = (x_ratio * indices[3]) - x0;; + + auto A = input(indices[0], indices[1], y0, x0); + auto B = input(indices[0], indices[1], y0, x1); + auto C = input(indices[0], indices[1], y1, x0); + auto D = input(indices[0], indices[1], y1, x1); + + return ((A*(cone-w)*(cone-h)) + (B*(w)*(cone-h)) + (C*(h)*(cone-w)) + (D*w*h)); + }, name, tag); +} + +/*! +* \brief Resize given tensor to given shape using bilinear interpolation +* +* \param input The input tensor. +* \param shape Output shape to resize to. +* \param layout input layout +* \param align_corners To preserve centers of 4 corner pixels +* \param name Name of the operation +* \param tag The tag to mark the operation +* +* \return A Tensor resized to given shape +*/ +inline Tensor resize_bilinear(const Tensor& input, + const Array& shape, + std::string layout = "NCHW", + bool align_corners = false, + std::string name = "tensor", + std::string tag = kInjective) { + Tensor ret; + + if (layout == "NHWC") { + ret = resize_bilinear_nhwc(input, shape, align_corners); + } else { + ret = resize_bilinear_nchw(input, shape, align_corners); + } + + return cast(ret, input->dtype); +} + +/*! +* \brief Resize given tensor to given shape +* +* \param input The input tensor. +* \param shape Output shape to resize to. +* \param layout input layout +* \param align_corners To preserve centers of 4 corner pixels +* \param mode Angorithm to use (NEAREST_NEIGHBOR / BILINEAR) +* \param name Name of the operation +* \param tag The tag to mark the operation +* +* \return A Tensor resized to given shape +*/ +inline Tensor resize(const Tensor& input, + const Array& shape, + std::string layout = "NCHW", + bool align_corners = false, + std::string mode = "BILINEAR", + std::string name = "tensor", + std::string tag = kInjective) { + if (mode == "NEAREST_NEIGHBOR") { + return resize_nearest_neighbor(input, shape, layout, align_corners); + } else { + return resize_bilinear(input, shape, layout, align_corners); + } +} + +} // namespace image +} // namespace topi +#endif // TOPI_IMAGE_RESIZE_H_ diff --git a/topi/include/topi/nn/upsampling.h b/topi/include/topi/nn/upsampling.h new file mode 100644 index 00000000..e0fe2991 --- /dev/null +++ b/topi/include/topi/nn/upsampling.h @@ -0,0 +1,44 @@ +/*! + * Copyright (c) 2017 by Contributors + * \file topi/nn/upsampling.h + * \brief upsampling op constructors + */ +#ifndef TOPI_NN_UPSAMPLING_H_ +#define TOPI_NN_UPSAMPLING_H_ + +#include +#include +#include +#include + +#include "topi/image/resize.h" + +namespace topi { +namespace nn { +using namespace tvm; +using namespace topi::image; + +/*! +* \brief Upsample given tensor to given shape +* +* \param input The input tensor. +* \param shape Output shape to upsample. +* \param layout input layout +* \param mode Angorithm to use (NEAREST_NEIGHBOR / BILINEAR) +* \param name Name of the operation +* \param tag The tag to mark the operation +* +* \return A Tensor upsampled to given shape +*/ +inline Tensor upsampling(const Tensor& input, + const Array shape, + std::string layout = "NCHW", + std::string mode = "NEAREST_NEIGHBOR", + std::string name = "tensor", + std::string tag = kInjective) { + return resize(input, shape, layout, false, mode); +} + +} // namespace nn +} // namespace topi +#endif // TOPI_NN_UPSAMPLING_H_ diff --git a/topi/python/topi/__init__.py b/topi/python/topi/__init__.py index f82a40f6..59549f89 100644 --- a/topi/python/topi/__init__.py +++ b/topi/python/topi/__init__.py @@ -30,6 +30,7 @@ from . import opengl from . import util from . import rocm from . import vision +from . import image # not import testing by default # because testing can have extra deps that are not necessary # we can import them from test cases explicitly diff --git a/topi/python/topi/cpp.py b/topi/python/topi/cpp.py index ae19484a..4c68dd73 100644 --- a/topi/python/topi/cpp.py +++ b/topi/python/topi/cpp.py @@ -50,6 +50,8 @@ vision = _create_module("vision") _init_api_prefix("topi.cpp.vision", "topi.vision") yolo2 = _create_module("vision.yolo2") _init_api_prefix("topi.cpp.vision.yolo2", "topi.vision.yolo2") +image = _create_module("image") +_init_api_prefix("topi.cpp.image", "topi.image") class IntVector(object): """Handle to std::vector instance """ diff --git a/topi/python/topi/image/__init__.py b/topi/python/topi/image/__init__.py new file mode 100644 index 00000000..840ae83e --- /dev/null +++ b/topi/python/topi/image/__init__.py @@ -0,0 +1,5 @@ +# pylint: disable=wildcard-import +"""IMAGE network operators""" +from __future__ import absolute_import as _abs + +from .resize import * diff --git a/topi/python/topi/image/resize.py b/topi/python/topi/image/resize.py new file mode 100644 index 00000000..461214cb --- /dev/null +++ b/topi/python/topi/image/resize.py @@ -0,0 +1,33 @@ +"""TVM operator input resize compute.""" +from __future__ import absolute_import +import topi + +def resize(data, size, layout="NCHW", align_corners=False, method="BILINEAR"): + """Perform resize operation on the data. + + Parameters + ---------- + inputs : tvm.Tensor + inputs is a 4-D tensor with shape + [batch, channel, in_height, in_width] + or [batch, in_height, in_width, channel] + + size: Tuple + Output resolution scale to + + layout: string, optional + either "NCHW" or "NHWC" + + align_corners: Boolean, optional + To preserve the values at the corner pixels + + method: {"BILINEAR", "NEAREST_NEIGHBOR"} + Method to be used for resizing. + + Returns + ------- + output : tvm.Tensor + 4-D with shape [batch, channel, in_height*scale, in_width*scale] + or [batch, in_height*scale, in_width*scale, channel] + """ + return topi.cpp.image.resize(data, size, layout, align_corners, method) diff --git a/topi/python/topi/nn/upsampling.py b/topi/python/topi/nn/upsampling.py index 9297eb4a..55f78443 100644 --- a/topi/python/topi/nn/upsampling.py +++ b/topi/python/topi/nn/upsampling.py @@ -1,25 +1,28 @@ """TVM operator upsampling compute.""" from __future__ import absolute_import -import tvm -from .. import util +import topi -def upsampling(data, scale, layout="NCHW"): - """Perform nearest neighbor upsampling on the data. - Bilinear upsampling is not supported. +def upsampling(data, scale, layout="NCHW", method='NEAREST_NEIGHBOR'): + """Perform upsampling on the data. + Nearest neighbor and bilinear upsampling are supported. Parameters ---------- - data : tvm.Tensor - 4-D with shape [batch, channel, in_height, in_width] + inputs : tvm.Tensor + inputs is a 4-D tensor with shape + [batch, channel, in_height, in_width] or [batch, in_height, in_width, channel] - scale: int - upsampling scaling factor + scale : int + Scaling factor - layout: string + layout : string, optional either "NCHW" or "NHWC" + method : {"BILINEAR", "NEAREST_NEIGHBOR"} + Method to be used for upsampling. + Returns ------- output : tvm.Tensor @@ -28,53 +31,10 @@ def upsampling(data, scale, layout="NCHW"): """ if layout == "NCHW": - return upsampling_nchw(data, scale) + out_shape = (data.shape[2] * scale, data.shape[3] * scale) elif layout == "NHWC": - return upsampling_nhwc(data, scale) + out_shape = (data.shape[1] * scale, data.shape[2] * scale) else: raise ValueError("not support this layout {} yet".format(layout)) - -def upsampling_nchw(data, scale): - """Perform nearest neighor upsampling on NCHW layout input. - - Parameters - ---------- - data : tvm.Tensor - 4-D with shape [batch, channel, in_height, in_width] - - scale: int - upsampling scaling factor - - Returns - ------- - output : tvm.Tensor - 4-D with shape [batch, channel, in_height*scale, in_width*scale] - """ - batch, channel, height, width = data.shape - out_height = util.simplify(height * scale) - out_width = util.simplify(width * scale) - - return tvm.compute((batch, channel, out_height, out_width), \ - lambda n, c, h, w: data[n, c, h/scale, w/scale]) - - -def upsampling_nhwc(data, scale): - """Perform nearest neighor upsampling on NHWC layout input. - - Parameters - ---------- - data : tvm.Tensor - 4-D with shape [batch, in_height, in_width, channel] - - scale: int - upsampling scaling factor - - """ - - batch, height, width, channel = data.shape - out_height = util.simplify(height * scale) - out_width = util.simplify(width * scale) - - return tvm.compute((batch, out_height, out_width, channel), \ - lambda n, h, w, c: data[n, h/scale, w/scale, c]) + return topi.cpp.nn.upsampling(data, out_shape, layout, method) diff --git a/topi/python/topi/testing/__init__.py b/topi/python/topi/testing/__init__.py index 66a86572..3731040e 100644 --- a/topi/python/topi/testing/__init__.py +++ b/topi/python/topi/testing/__init__.py @@ -12,6 +12,7 @@ from .depthwise_conv2d_python import depthwise_conv2d_python_nchw, depthwise_con from .dilate_python import dilate_python from .softmax_python import softmax_python, log_softmax_python from .upsampling_python import upsampling_python +from .bilinear_resize_python import bilinear_resize_python from .reorg_python import reorg_python from .region_python import region_python from .shortcut_python import shortcut_python diff --git a/topi/python/topi/testing/bilinear_resize_python.py b/topi/python/topi/testing/bilinear_resize_python.py new file mode 100644 index 00000000..de931d7e --- /dev/null +++ b/topi/python/topi/testing/bilinear_resize_python.py @@ -0,0 +1,79 @@ +# pylint: disable=invalid-name, line-too-long, unused-variable, too-many-locals +"""Bilinear Scale in python""" +import math +import numpy as np + +def bilinear_weights(height, width, new_h, new_w, align_corners=False): + """ Helper function to generate weights for bilinear scaling """ + + if align_corners: + x_ratio = np.float32(np.float32(width)/np.float32(new_w)) + y_ratio = np.float32(np.float32(height)/np.float32(new_h)) + else: + x_ratio = np.float32(np.float32(width-1)/np.float32(new_w-1)) + y_ratio = np.float32(np.float32(height-1)/np.float32(new_h-1)) + + def _bilinear_interpolation(y, x): + x_coord = math.floor(x_ratio * x) + y_coord = math.floor(y_ratio * y) + x_diff = np.float32((x_ratio * x) - x_coord) + y_diff = np.float32((y_ratio * y) - y_coord) + + return [y_coord, x_coord, y_diff, x_diff] + + # weights to hold (srcx, srcy, x_diff, y_diff) for each out value. + weights = np.empty([new_h, new_w, 4], dtype='float32') + + for i in range(new_h): + for j in range(new_w): + weights[i][j] = _bilinear_interpolation(i, j) + return weights + +def bilinear_resize_python(image, out_size, layout, align_corners=False): + """ Bilinear scaling using python""" + (new_h, new_w) = out_size + + if layout == 'NHWC': + (batch, h, w, channel) = image.shape + scaled_image = np.ones((batch, new_h, new_w, channel)) + else: + (batch, channel, h, w) = image.shape + scaled_image = np.ones((batch, channel, new_h, new_w)) + + weights = bilinear_weights(h, w, new_h, new_w, align_corners) + + for b in range(batch): + for i in range(channel): + for j in range(new_h): + for k in range(new_w): + y0 = int(weights[j][k][0]) + x0 = int(weights[j][k][1]) + + x1 = min((x0+1), (w-1)) + y1 = min((y0+1), (h-1)) + + y_diff = weights[j][k][2] + x_diff = weights[j][k][3] + + if layout == 'NHWC': + A = image[b][y0][x0][i] + B = image[b][y0][x1][i] + C = image[b][y1][x0][i] + D = image[b][y1][x1][i] + else: + A = image[b][i][y0][x0] + B = image[b][i][y0][x1] + C = image[b][i][y1][x0] + D = image[b][i][y1][x1] + + pixel = np.float32((A*(1-x_diff)*(1-y_diff) + + B*(x_diff)*(1-y_diff) + + C*(y_diff)*(1-x_diff) + + D*(x_diff*y_diff))) + + if layout == 'NHWC': + scaled_image[b][j][k][i] = pixel + else: + scaled_image[b][i][j][k] = pixel + + return scaled_image diff --git a/topi/python/topi/testing/upsampling_python.py b/topi/python/topi/testing/upsampling_python.py index 328c7a5a..fc4ad652 100644 --- a/topi/python/topi/testing/upsampling_python.py +++ b/topi/python/topi/testing/upsampling_python.py @@ -3,13 +3,26 @@ import numpy as np def upsample_nearest(arr, scale): + """ Populate the array by scale factor""" return arr.repeat(scale, axis=0).repeat(scale, axis=1) -def upsampling_python(data, scale): +def upsampling_python(data, scale, layout='NCHW'): + """ Python version of scaling using nearest neighbour """ + ishape = data.shape - oshape = (ishape[0], ishape[1], ishape[2]*scale, ishape[3]*scale) - output_np = np.zeros(oshape, dtype=data.dtype) - for b in range(oshape[0]): - for c in range(oshape[1]): - output_np[b, c, :, :] = upsample_nearest(data[b, c, :, :], scale) - return output_np + if layout == 'NCHW': + oshape = (ishape[0], ishape[1], ishape[2]*scale, ishape[3]*scale) + output_np = np.zeros(oshape, dtype=data.dtype) + for b in range(oshape[0]): + for c in range(oshape[1]): + output_np[b, c, :, :] = upsample_nearest(data[b, c, :, :], scale) + return output_np + elif layout == 'NHWC': + oshape = (ishape[0], ishape[1]*scale, ishape[1]*scale, ishape[3]) + output_np = np.zeros(oshape, dtype=data.dtype) + for b in range(oshape[0]): + for c in range(oshape[3]): + output_np[b, :, :, c] = upsample_nearest(data[b, :, :, c], scale) + return output_np + else: + raise ValueError("not support this layout {} yet".format(layout)) diff --git a/topi/src/topi.cc b/topi/src/topi.cc index 9c820edf..93301d70 100644 --- a/topi/src/topi.cc +++ b/topi/src/topi.cc @@ -23,8 +23,10 @@ #include #include #include +#include #include +#include #include #include #include @@ -285,6 +287,12 @@ TVM_REGISTER_GLOBAL("topi.strided_slice") *rv = strided_slice(args[0], args[1], args[2], args[3]); }); +/* Ops from nn/upsampling.h */ +TVM_REGISTER_GLOBAL("topi.nn.upsampling") +.set_body([](TVMArgs args, TVMRetValue *rv) { + *rv = nn::upsampling(args[0], args[1], args[2], args[3]); + }); + /* Ops from nn/batch_norm.h */ TVM_REGISTER_GLOBAL("topi.nn.batch_norm_inference") .set_body([](TVMArgs args, TVMRetValue *rv) { @@ -366,10 +374,18 @@ TVM_REGISTER_GLOBAL("topi.vision.reorg") .set_body([](TVMArgs args, TVMRetValue *rv) { *rv = vision::reorg(args[0], args[1]); }); + TVM_REGISTER_GLOBAL("topi.vision.yolo2.region") .set_body([](TVMArgs args, TVMRetValue *rv) { *rv = vision::yolo2::region(args[0], args[1], args[2], args[3], args[4], args[5]); }); + +/* Ops from image/resize.h */ +TVM_REGISTER_GLOBAL("topi.image.resize") +.set_body([](TVMArgs args, TVMRetValue *rv) { + *rv = image::resize(args[0], args[1], args[2], args[3], args[4]); + }); + /* Generic schedules */ TVM_REGISTER_GLOBAL("topi.generic.default_schedule") .set_body([](TVMArgs args, TVMRetValue *rv) { diff --git a/topi/tests/python/test_topi_resize.py b/topi/tests/python/test_topi_resize.py new file mode 100644 index 00000000..cb83c18d --- /dev/null +++ b/topi/tests/python/test_topi_resize.py @@ -0,0 +1,57 @@ +"""Test code for bilinear scale """ +import numpy as np +import tvm +import topi +import topi.testing +import math + +def verify_bilinear_scale(batch, in_channel, in_height, in_width, out_height, out_width, layout='NCHW', align_corners=False): + + if layout == 'NCHW': + A = tvm.placeholder((batch, in_channel, in_height, in_width), name='A', dtype='float32') + dtype = A.dtype + out_shape = (batch, in_channel, out_height, out_width) + a_np = np.random.uniform(size=(batch, in_channel, in_height, in_width)).astype(dtype) + elif layout == 'NHWC': + A = tvm.placeholder((batch, in_height, in_width, in_channel), name='A', dtype='float32') + dtype = A.dtype + out_shape = (batch, out_height, out_width, in_channel) + a_np = np.random.uniform(size=(batch, in_height, in_width, in_channel)).astype(dtype) + else: + raise NotImplementedError( + 'Layout not supported {} '.format(layout)) + + B = topi.image.resize(A, (out_height, out_width), layout=layout, align_corners=align_corners) + + b_np = topi.testing.bilinear_resize_python(a_np, (out_height, out_width), layout, align_corners) + + 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) + + np.testing.assert_allclose(b.asnumpy(), b_np, rtol=1e-3, atol=1e-3) + + for device in ['llvm', 'cuda', 'vulkan']: + check_device(device) + +def test_resize(): + # Scale NCHW + verify_bilinear_scale(4, 16, 32, 32, 50, 50, 'NCHW') + # Scale NCHW + Align Corners + verify_bilinear_scale(6, 32, 64, 64, 20, 20, 'NCHW', True) + # Scale NHWC + verify_bilinear_scale(4, 16, 32, 32, 50, 50, "NHWC") + # Scale NHWC + Align Corners + verify_bilinear_scale(6, 32, 64, 64, 20, 20, "NHWC", True) + +if __name__ == "__main__": + test_resize() diff --git a/topi/tests/python/test_topi_upsampling.py b/topi/tests/python/test_topi_upsampling.py index 7421dd41..d03ea0bc 100644 --- a/topi/tests/python/test_topi_upsampling.py +++ b/topi/tests/python/test_topi_upsampling.py @@ -5,14 +5,26 @@ import topi import topi.testing import math -def verify_upsampling(batch, in_channel, in_height, in_width, scale): - A = tvm.placeholder((batch, in_channel, in_height, in_width), name='A') - B = topi.nn.upsampling(A, scale) - out_shape = (batch, in_channel, in_height*scale, in_width*scale) - dtype = A.dtype +def verify_upsampling(batch, in_channel, in_height, in_width, scale, layout='NCHW'): - a_np = np.random.uniform(size=(batch, in_channel, in_height, in_width)).astype(dtype) - b_np = topi.testing.upsampling_python(a_np, scale) + + if layout == 'NCHW': + A = tvm.placeholder((batch, in_channel, in_height, in_width), name='A') + dtype = A.dtype + out_shape = (batch, in_channel, in_height*scale, in_width*scale) + a_np = np.random.uniform(size=(batch, in_channel, in_height, in_width)).astype(dtype) + elif layout == 'NHWC': + A = tvm.placeholder((batch, in_height, in_width, in_channel), name='A') + dtype = A.dtype + out_shape = (batch, in_height*scale, in_width*scale, in_channel) + a_np = np.random.uniform(size=(batch, in_height, in_width, in_channel)).astype(dtype) + else: + raise NotImplementedError( + 'Layout not supported {} '.format(layout)) + + B = topi.nn.upsampling(A, scale, layout=layout) + + b_np = topi.testing.upsampling_python(a_np, scale, layout) def check_device(device): ctx = tvm.context(device, 0) @@ -33,8 +45,12 @@ def verify_upsampling(batch, in_channel, in_height, in_width, scale): check_device(device) def test_upsampling(): + # NCHW verify_upsampling(8, 16, 32, 32, 2) verify_upsampling(12, 32, 64, 64, 3) + # NHWC + verify_upsampling(8, 16, 32, 32, 2, "NHWC") + verify_upsampling(12, 32, 64, 64, 3, "NHWC") if __name__ == "__main__": test_upsampling()