[Relay][Op] Adaptive pooling

include/tvm/relay/attrs/nn.h

@@ -332,6 +332,22 @@ struct GlobalPool2DAttrs : public tvm::AttrsNode<GlobalPool2DAttrs> {
   }
 };
 
+/*! \brief Attributes for adaptive pool operator */
+struct AdaptivePool2DAttrs : public tvm::AttrsNode<AdaptivePool2DAttrs> {
+  Array<IndexExpr> output_size;
+  std::string layout;
+
+  TVM_DECLARE_ATTRS(AdaptivePool2DAttrs, "relay.attrs.AdaptivePool2DAttrs") {
+    TVM_ATTR_FIELD(output_size).set_default(Array<IndexExpr>({}))
+      .describe("Output height and width.");
+    TVM_ATTR_FIELD(layout).set_default("NCHW")
+      .describe("Dimension ordering of data and weight. 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"
+                  "'W' dimensions.");
+  }
+};
+
 
 /*! \brief Attributes for dense operator */
 struct DenseAttrs : public tvm::AttrsNode<DenseAttrs> {

nnvm/python/nnvm/top/nn.py

@@ -399,7 +399,7 @@ def schedule_avg_pool2d(attrs, outs, target):
 def schedule_global_max_pool2d(_, outs, target):
     """Schedule definition of global_max_pool2d"""
     with tvm.target.create(target):
-        return topi.generic.schedule_global_pool(outs)
+        return topi.generic.schedule_adaptive_pool(outs)
 
 reg.register_pattern("global_max_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 
@@ -409,7 +409,7 @@ def schedule_global_max_pool2d(_, outs, target):
 def schedule_global_avg_pool2d(_, outs, target):
     """Schedule definition of global_avg_pool2d"""
     with tvm.target.create(target):
-        return topi.generic.schedule_global_pool(outs)
+        return topi.generic.schedule_adaptive_pool(outs)
 
 reg.register_pattern("global_avg_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 

python/tvm/relay/frontend/mxnet.py

@@ -190,10 +190,7 @@ def _pool2d(new_op, is_avg):
 
 def _mx_adaptive_avg_pooling(inputs, attrs):
     output_size = attrs.get_int_tuple("output_size", [])
-    if output_size != (1,):
-        raise tvm.error.OpAttributeUnimplemented(
-            "AdaptiveAvgPooling with output_size other than 1 is not supported yet.")
-    return _op.nn.global_avg_pool2d(inputs[0])
+    return _op.nn.contrib_adaptive_avg_pool2d(inputs[0], output_size)
 
 
 def _mx_dropout(inputs, attrs):

python/tvm/relay/op/nn/_nn.py

@@ -247,7 +247,7 @@ def schedule_avg_pool2d(attrs, outs, target):
 def schedule_global_max_pool2d(_, outs, target):
     """Schedule definition of global_max_pool2d"""
     with target:
-        return topi.generic.schedule_global_pool(outs)
+        return topi.generic.schedule_adaptive_pool(outs)
 
 
 reg.register_pattern("nn.global_max_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
@@ -258,11 +258,32 @@ def schedule_global_max_pool2d(_, outs, target):
 def schedule_global_avg_pool2d(_, outs, target):
     """Schedule definition of global_avg_pool2d"""
     with target:
-        return topi.generic.schedule_global_pool(outs)
+        return topi.generic.schedule_adaptive_pool(outs)
 
 
 reg.register_pattern("nn.global_avg_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 
+
+# adaptive_max_pool2d
+@reg.register_schedule("nn.contrib_adaptive_max_pool2d")
+def schedule_contrib_adaptive_max_pool2d(_, outs, target):
+    """Schedule definition of adaptive_max_pool2d"""
+    with target:
+        return topi.generic.schedule_adaptive_pool(outs)
+
+reg.register_pattern("nn.contrib_adaptive_max_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
+
+
+# adaptive_avg_pool2d
+@reg.register_schedule("nn.contrib_adaptive_avg_pool2d")
+def schedule_contrib_adaptive_avg_pool2d(_, outs, target):
+    """Schedule definition of adaptive_avg_pool2d"""
+    with target:
+        return topi.generic.schedule_adaptive_pool(outs)
+
+reg.register_pattern("nn.contrib_adaptive_avg_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
+
+
 # leaky_relu
 reg.register_schedule("nn.leaky_relu", schedule_broadcast)
 reg.register_pattern("nn.leaky_relu", OpPattern.ELEMWISE)

python/tvm/relay/op/nn/nn.py

@@ -390,6 +390,99 @@ def global_avg_pool2d(data,
     return _make.global_avg_pool2d(data, layout)
 
 
+def contrib_adaptive_max_pool2d(data,
+                                output_size=None,
+                                layout="NCHW"):
+    r"""2D adaptive max pooling operator. This operator is experimental.
+
+    This operator takes data as input and does 2D max value calculation
+    across each window represented by WxH.
+
+
+    In the default case, where the data_layout is `NCHW`
+    a data Tensor with shape `(batch_size, in_channels, height, width)`,
+    to produce an output Tensor with shape
+    (batch_size, in_channels, output_height, output_width).
+
+    The pooling kernel and stride sizes are automatically chosen for
+    desired output sizes.
+
+    For output_size:
+        If this argument is not provided, input height and width will be used
+        as output height and width.
+
+        If a single integer is provided for output_size, the output size is
+        (N x C x output_size x output_size) for any input (NCHW).
+
+        If a tuple of integers (height, width) are provided for output_size,
+        the output size is (N x C x height x width) for any input (NCHW).
+
+    Parameters
+    ----------
+    data : tvm.relay.Expr
+        The input data to the operator.
+
+    output_size : tuple of int. optional
+        Output height and width.
+
+    layout : str, optional
+        Layout of the input.
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The computed result.
+    """
+    output_size = [] or output_size
+    return _make.contrib_adaptive_max_pool2d(data, output_size, layout)
+
+def contrib_adaptive_avg_pool2d(data,
+                                output_size=None,
+                                layout="NCHW"):
+    r"""2D adaptive average pooling operator. This operator is experimental.
+
+    This operator takes data as input and does 2D average value calculation
+    across each window represented by WxH.
+
+
+    In the default case, where the data_layout is `NCHW`
+    a data Tensor with shape `(batch_size, in_channels, height, width)`,
+    to produce an output Tensor with shape
+    (batch_size, in_channels, output_height, output_width).
+
+    The pooling kernel and stride sizes are automatically chosen for
+    desired output sizes.
+
+    For output_size:
+        If this argument is not provided, input height and width will be used
+        as output height and width.
+
+        If a single integer is provided for output_size, the output size is
+        (N x C x output_size x output_size) for any input (NCHW).
+
+        If a tuple of integers (height, width) are provided for output_size,
+        the output size is (N x C x height x width) for any input (NCHW).
+
+    Parameters
+    ----------
+    data : tvm.relay.Expr
+        The input data to the operator.
+
+    output_size : tuple of int. optional
+        Output height and width.
+
+    layout : str, optional
+        Layout of the input.
+
+    Returns
+    -------
+    result : tvm.relay.Expr
+        The computed result.
+    """
+    output_size = [] or output_size
+    return _make.contrib_adaptive_avg_pool2d(data, output_size, layout)
+
+
 def upsampling(data,
                scale=1,
                layout="NCHW",

src/relay/op/nn/pooling.cc

@@ -72,7 +72,6 @@ bool Pool2DRel(const Array<Type>& types,
 
   CHECK(data != nullptr);
   const auto dshape = data->shape;
-  CHECK_NE(dshape.size(), 0);
   CHECK_GE(dshape.size(), 2U)
       << "Pool2D only support input >= 2-D: input must have height and width";
   const auto param = attrs.as<AttrType>();
@@ -284,7 +283,6 @@ bool GlobalPool2DRel(const Array<Type>& types,
   const auto* data = types[0].as<TensorTypeNode>();
   if (data == nullptr) { return false; }
   const auto dshape = data->shape;
-  CHECK_NE(dshape.size(), 0);
   CHECK_GE(dshape.size(), 2U)
       << "Pool2D only support input >= 2-D: input must have height and width";
   const auto param = attrs.as<GlobalPool2DAttrs>();
@@ -393,5 +391,170 @@ RELAY_REGISTER_OP("nn.global_max_pool2d")
                                Pool2DInferCorrectLayout<GlobalPool2DAttrs>)
 .set_attr<FTVMCompute>("FTVMCompute", GlobalPool2DCompute<topi::nn::kMaxPool>);
 
+
+// relay.nn.adaptive_pool_2d
+TVM_REGISTER_NODE_TYPE(AdaptivePool2DAttrs);
+
+bool AdaptivePool2DRel(const Array<Type>& types,
+                       int num_inputs,
+                       const Attrs& attrs,
+                       const TypeReporter& reporter) {
+  CHECK_EQ(types.size(), 2);
+  const auto* data = types[0].as<TensorTypeNode>();
+  if (data == nullptr) { return false; }
+  const auto dshape = data->shape;
+  CHECK_GE(dshape.size(), 2U)
+    << "Pool2D only support input >= 2-D: input must have height and width";
+  const auto* param = attrs.as<AdaptivePool2DAttrs>();
+  CHECK(param != nullptr);
+
+  Layout layout(param->layout);
+  CHECK(layout.Contains(LayoutAxis::Get('H')) && layout.Contains(LayoutAxis::Get('W')) &&
+        !layout.Contains(LayoutAxis::Get('h')) && !layout.Contains(LayoutAxis::Get('w')))
+    << "Invalid layout " << layout
+    << ". Pool2D layout must have H and W, which cannot be split";
+
+  const auto hidx = layout.IndexOf(LayoutAxis::Get('H'));
+  const auto widx = layout.IndexOf(LayoutAxis::Get('W'));
+  Array<IndexExpr> oshape(dshape);
+  auto output_size = param->output_size;
+  CHECK_LE(output_size.size(), 2U)
+    << "output_size can have up to 2 elements.";
+  IndexExpr output_height, output_width;
+  if (output_size.empty()) {
+    output_height = dshape[hidx];
+    output_width = dshape[widx];
+  } else if (output_size.size() == 1) {
+    output_height = output_size[0];
+    output_width = output_size[0];
+  } else {
+    output_height = output_size[0];
+    output_width = output_size[1];
+  }
+
+  oshape.Set(hidx, output_height);
+  oshape.Set(widx, output_width);
+
+  // assign output type
+  reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype));
+  return true;
+}
+
+template<topi::nn::PoolType mode>
+Array<Tensor> AdaptivePool2DCompute(const Attrs& attrs,
+                                    const Array<Tensor>& inputs,
+                                    const Type& out_type,
+                                    const Target& target) {
+  static const Layout kNCHW("NCHW");
+  const auto* param = attrs.as<AdaptivePool2DAttrs>();
+  CHECK(param != nullptr);
+  Layout layout(param->layout);
+  CHECK(BijectiveLayoutNode::make(layout, kNCHW).defined())
+    << "Adaptive pool2d currently only supports layouts that are convertible from NCHW";
+  CHECK_EQ(layout.IndexOf(LayoutAxis::Get('h')), -1)
+    << "Adaptive pool2d does not support input split on height";
+  CHECK_EQ(layout.IndexOf(LayoutAxis::Get('w')), -1)
+    << "Adaptive pool2d does not support input split on width";
+
+  CHECK(inputs[0].ndim() == 4U || inputs[0].ndim() == 5U)
+    << "Pool2D only support 4-D input (e.g., NCHW)"
+    << " or 5-D input (last dimension is a split of channel)";
+
+  auto output_size = param->output_size;
+  const auto hidx = layout.IndexOf(LayoutAxis::Get('H'));
+  const auto widx = layout.IndexOf(LayoutAxis::Get('W'));
+  IndexExpr output_height, output_width;
+  if (output_size.empty()) {
+    output_height = inputs[0]->shape[hidx];
+    output_width = inputs[0]->shape[widx];
+  } else if (output_size.size() == 1) {
+    output_height = output_size[0];
+    output_width = output_size[0];
+  } else {
+    output_height = output_size[0];
+    output_width = output_size[1];
+  }
+  return Array<Tensor>{
+    topi::nn::adaptive_pool(inputs[0], Array<IndexExpr>{ output_height, output_width },
+                            mode, layout.name()) };
+}
+
+// relay.nn.adaptive_avg_pool2d
+Expr MakeAdaptiveAvgPool2D(Expr data,
+                           Array<IndexExpr> output_size,
+                           std::string layout) {
+  auto attrs = make_node<AdaptivePool2DAttrs>();
+  attrs->output_size = std::move(output_size);
+  attrs->layout = std::move(layout);
+  static const Op& op = Op::Get("nn.contrib_adaptive_avg_pool2d");
+  return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_API("relay.op.nn._make.contrib_adaptive_avg_pool2d")
+.set_body_typed(MakeAdaptiveAvgPool2D);
+
+RELAY_REGISTER_OP("nn.contrib_adaptive_avg_pool2d")
+  .describe(R"code(Adaptive average pooling operation for 2D data.
+
+- **data**: This depends on the `layout` parameter. Input is 4D array of shape
+            (batch_size, channels, height, width) if `layout` is `NCHW`.
+- **output_size**: If this argument is not provided, input height and width will be used
+                   as output height and width.
+                   If a single integer is provided for output_size, the output size is
+                   (N x C x output_size x output_size) for any input (NCHW).
+                   If a tuple of integers (height, width) are provided for output_size,
+                   the output size is (N x C x height x width) for any input (NCHW).
+- **out**: This depends on the `layout` parameter. Output is 4D array of shape
+           (batch_size, channels, output_height, output_width)  if `layout` is `NCHW`.
+
+)code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.AdaptivePool2DAttrs")
+.set_num_inputs(1)
+.add_argument("data", "Tensor", "The input tensor.")
+.set_support_level(10)
+.add_type_rel("AdaptiveAvgPool2D", AdaptivePool2DRel)
+.set_attr<FInferCorrectLayout>("FInferCorrectLayout",
+                               Pool2DInferCorrectLayout<AdaptivePool2DAttrs>)
+.set_attr<FTVMCompute>("FTVMCompute", AdaptivePool2DCompute<topi::nn::kAvgPool>);
+
+
+// relay.nn.adaptive_max_pool2d
+Expr MakeAdaptiveMaxPool2D(Expr data,
+                           Array<IndexExpr> output_size,
+                           std::string layout) {
+  auto attrs = make_node<AdaptivePool2DAttrs>();
+  attrs->output_size = std::move(output_size);
+  attrs->layout = std::move(layout);
+  static const Op& op = Op::Get("nn.contrib_adaptive_max_pool2d");
+  return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_API("relay.op.nn._make.contrib_adaptive_max_pool2d")
+.set_body_typed(MakeAdaptiveMaxPool2D);
+
+RELAY_REGISTER_OP("nn.contrib_adaptive_max_pool2d")
+  .describe(R"code(Adaptive max pooling operation for 2D data.
+
+- **data**: This depends on the `layout` parameter. Input is 4D array of shape
+            (batch_size, channels, height, width) if `layout` is `NCHW`.
+- **output_size**: If this argument is not provided, input height and width will be used
+                   as output height and width.
+                   If a single integer is provided for output_size, the output size is
+                   (N x C x output_size x output_size) for any input (NCHW).
+                   If a tuple of integers (height, width) are provided for output_size,
+                   the output size is (N x C x height x width) for any input (NCHW).
+- **out**: This depends on the `layout` parameter. Output is 4D array of shape
+           (batch_size, channels, output_height, output_width)  if `layout` is `NCHW`.
+
+)code" TVM_ADD_FILELINE)
+.set_attrs_type_key("relay.attrs.AdaptivePool2DAttrs")
+.set_num_inputs(1)
+.add_argument("data", "Tensor", "The input tensor.")
+.set_support_level(10)
+.add_type_rel("AdaptiveMaxPool2D", AdaptivePool2DRel)
+.set_attr<FInferCorrectLayout>("FInferCorrectLayout",
+                               Pool2DInferCorrectLayout<AdaptivePool2DAttrs>)
+.set_attr<FTVMCompute>("FTVMCompute", AdaptivePool2DCompute<topi::nn::kMaxPool>);
+
 }  // namespace relay
 }  // namespace tvm