[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> {

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

@@ -263,6 +263,27 @@ def schedule_global_avg_pool2d(_, outs, target):
 
 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

tests/python/relay/test_op_level10.py

@@ -208,7 +208,50 @@ def test_shape_of():
             tvm.testing.assert_allclose(op_res.asnumpy(),
                                         np.array(shape).astype('int32'))
 
+def verify_adaptive_pool2d(dshape, out_size, pool_type, layout="NCHW", dtype="float32"):
+    def start_index(index, odim, idim):
+        return int(np.floor(index * idim / odim))
+
+    def end_index(index, odim, idim):
+        return int(np.ceil((index + 1) * idim / odim))
+
+    np_data = np.random.uniform(low=0, high=255, size=dshape).astype(dtype)
+    n, c, h, w = dshape
+    oh, ow = out_size
+    oshape = (n, c) + out_size
+    np_out = np.zeros(oshape).astype(dtype)
+    np_op = np.mean if pool_type == "avg" else np.max
+    for i in range(n):
+        for j in range(c):
+            for k in range(oh):
+                k_start = start_index(k, oh, h)
+                k_end = end_index(k, oh, h)
+                k_sl = slice(k_start, k_end)
+                for l in range(ow):
+                    l_start = start_index(l, ow, w)
+                    l_end = end_index(l, ow, w)
+                    l_sl = slice(l_start, l_end)
+                    np_out[i, j, k, l] = np_op(np_data[i, j, k_sl, l_sl])
+
+    opfunc = relay.nn.contrib_adaptive_avg_pool2d if pool_type == "avg" else relay.nn.contrib_adaptive_max_pool2d
+    x = relay.var("x", relay.TensorType((n, c, h, w), "float32"))
+    y = opfunc(x, out_size, layout)
+    func = relay.Function([x], y)
+
+    for target, ctx in ctx_list():
+        intrp1 = relay.create_executor("graph", ctx=ctx, target=target)
+        relay_out = intrp1.evaluate(func)(np_data)
+        tvm.testing.assert_allclose(relay_out.asnumpy(), np_out, rtol=1e-5, atol=1e-5)
+
+def test_adaptive_pool2d():
+    verify_adaptive_pool2d((1, 9, 224, 224), (1, 1), "max")
+    verify_adaptive_pool2d((1, 3, 224, 224), (2, 3), "avg")
+    verify_adaptive_pool2d((1, 14, 56, 78), (34, 13), "max")
+    verify_adaptive_pool2d((1, 5, 46, 97), (4, 96), "avg")
+
+
 if __name__ == "__main__":
+    test_adaptive_pool2d()
     test_collapse_sum_like()
     test_broadcast_to_like()
     test_slice_like()

tests/python/relay/test_op_level2.py

@@ -316,7 +316,6 @@ def test_avg_pool2d_no_count_pad():
         op_res1 = intrp1.evaluate(func)(data)
         tvm.testing.assert_allclose(op_res1.asnumpy(), ref_res, rtol=1e-5, atol=1e-5)
 
-
 def test_flatten_infer_type():
     d1, d2, d3, d4 = tvm.var("d1"), tvm.var("d2"), tvm.var("d3"), tvm.var("d4")
     x = relay.var("x", relay.TensorType((d1, d2, d3, d4), "float32"))