[TOPI][IMAGE][RESIZE] Bilinear interpolation for resize and upsampling.

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
 ~~~~~~~~~~~~

nnvm/include/nnvm/top/nn.h

@@ -288,16 +288,22 @@ struct GlobalPool2DParam : public dmlc::Parameter<GlobalPool2DParam> {
 struct UpSamplingParam : public dmlc::Parameter<UpSamplingParam> {
   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");
   }
 };
 

nnvm/python/nnvm/top/__init__.py

@@ -8,6 +8,7 @@
 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

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)

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)
 

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 <tvm/tvm.h>
+#include <tvm/expr.h>
+#include <tvm/packed_func_ext.h>
+#include <nnvm/layout.h>
+#include <nnvm/compiler/op_attr_types.h>
+#include <nnvm/op.h>
+#include <nnvm/node.h>
+#include <nnvm/op_attr_types.h>
+#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<TShape>* in_shape,
+                             std::vector<TShape>* out_shape) {
+  static const Layout kNCHW("NCHW");
+  const ResizeParam& param = nnvm::get<ResizeParam>(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<Layout> *in_layouts,
+                         const std::vector<Layout> *last_in_layouts,
+                         std::vector<Layout> *out_layouts) {
+  const ResizeParam& param = nnvm::get<ResizeParam>(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<ResizeParam>)
+.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<ResizeParam>)
+.set_attr<FInferShape>("FInferShape", ResizeInferShape)
+.set_attr<FInferType>("FInferType", ElemwiseType<1, 1>)
+.set_attr<FCorrectLayout>("FCorrectLayout", ResizeLayout)
+.set_num_outputs(1)
+.set_num_inputs(1)
+.set_attr<FTVMCompute>(
+  "FTVMCompute", [](const NodeAttrs& attrs,
+                    const Array<Tensor>& inputs,
+                    const Array<Tensor>& out_info) {
+  const ResizeParam& param = nnvm::get<ResizeParam>(attrs.parsed);
+  Array<Expr> 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<Tensor>{ topi::image::resize(inputs[0], oshape, param.layout,
+                                             param.align_corners, param.method)};
+})
+.set_support_level(2);
+
+}  // namespace top
+}  // namespace nnvm

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 <string>
+#include <vector>
+#include <utility>
+#include <iostream>
+#include <sstream>
+
+namespace nnvm {
+namespace top {
+
+struct ResizeParam : public dmlc::Parameter<ResizeParam> {
+  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_

nnvm/src/top/nn/upsampling.cc

@@ -1,36 +1,49 @@
 /*!
  *  Copyright (c) 2017 by Contributors
- * \file pooling.cc
- * \brief Property def of pooling operators.
+ * \file upsampling.cc
+ * \brief Property def of upsampling operators.
  */
+#include <tvm/tvm.h>
+#include <tvm/expr.h>
+#include <nnvm/layout.h>
+#include <nnvm/compiler/op_attr_types.h>
 #include <nnvm/op.h>
 #include <nnvm/node.h>
 #include <nnvm/op_attr_types.h>
 #include <nnvm/top/nn.h>
 #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<TShape>* in_shape,
-                                   std::vector<TShape>* out_shape) {
+                                 std::vector<TShape>* in_shape,
+                                 std::vector<TShape>* out_shape) {
   static const Layout kNCHW("NCHW");
   const UpSamplingParam& param = nnvm::get<UpSamplingParam>(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,13 +61,22 @@ 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**: data is 4D array of shape
+            (batch_size, channels, in_height, in_width) for NCHW
+            (batch_size, in_height, in_width, channels) for NHWC
 
-- **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).
+- **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.")
+.add_argument("weight", "3D Tensor", "Weight matrix.")
 .add_arguments(UpSamplingParam::__FIELDS__())
 .set_attr_parser(ParamParser<UpSamplingParam>)
 .set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<UpSamplingParam>)
@@ -63,6 +85,22 @@ NNVM_REGISTER_OP(upsampling)
 .set_attr<FCorrectLayout>("FCorrectLayout", UpsamplingLayout)
 .set_num_outputs(1)
 .set_num_inputs(1)
+.set_attr<FTVMCompute>(
+  "FTVMCompute", [](const NodeAttrs& attrs,
+                    const Array<Tensor>& inputs,
+                    const Array<Tensor>& out_info) {
+  const UpSamplingParam& param = nnvm::get<UpSamplingParam>(attrs.parsed);
+  Array<Expr> 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<Tensor>{ topi::nn::upsampling(inputs[0], oshape, param.layout, param.method)};
+})
 .set_support_level(2);
 
 }  // namespace top