[Relay][Op] Adaptive pooling (#3085)

* Add topi adaptive_pool

* Use adaptive_pool to compute global_pool

* Add relay adaptive pool2d

* Fix lint

* Fix typo

* Minor change

* Change support level to 10

* Add contrib

* Remove global pool schedule

* Add contrib module

* Fix lint

* Update doc

* Update doc

docs/api/python/topi.rst

@@ -57,6 +57,7 @@ List of operators
    topi.nn.dilate
    topi.nn.pool
    topi.nn.global_pool
+   topi.nn.adaptive_pool
    topi.nn.upsampling
    topi.nn.softmax
    topi.nn.dense

docs/langref/relay_op.rst

@@ -189,6 +189,8 @@ This level support backpropagation of broadcast operators. It is temporary.
    tvm.relay.annotation.on_device
    tvm.relay.reverse_reshape
    tvm.relay.nn.batch_matmul
+   tvm.relay.contrib.adaptive_max_pool2d
+   tvm.relay.contrib.adaptive_avg_pool2d
 
 
 Level 1 Definitions
@@ -318,3 +320,5 @@ Level 10 Definitions
 .. autofunction:: tvm.relay.annotation.on_device
 .. autofunction:: tvm.relay.reverse_reshape
 .. autofunction:: tvm.relay.nn.batch_matmul
+.. autofunction:: tvm.relay.contrib.adaptive_max_pool2d
+.. autofunction:: tvm.relay.contrib.adaptive_avg_pool2d

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 @@ reg.register_pattern("avg_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 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 @@ reg.register_pattern("global_max_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 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/__init__.py

@@ -40,6 +40,7 @@ from .op.algorithm import *
 from . import nn
 from . import annotation
 from . import vision
+from . import contrib
 from . import image
 from . import frontend
 from . import backend

python/tvm/relay/contrib.py

+# 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=wildcard-import, unused-import, unused-wildcard-import
+"""Contrib operators."""
+# Re-export in a specific file name so that autodoc can pick it up
+from .op.contrib import *

python/tvm/relay/frontend/mxnet.py

@@ -190,10 +190,7 @@ def _mx_pooling(inputs, attrs):
 
 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.contrib.adaptive_avg_pool2d(inputs[0], output_size)
 
 
 def _mx_dropout(inputs, attrs):

python/tvm/relay/op/__init__.py

@@ -29,6 +29,7 @@ from . import nn
 from . import annotation
 from . import image
 from . import vision
+from . import contrib
 from . import op_attrs
 
 

python/tvm/relay/op/contrib/__init__.py

+# 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=wildcard-import
+"""Neural network related operators."""
+from __future__ import absolute_import as _abs
+from .contrib import *
+from . import _contrib

python/tvm/relay/op/contrib/_contrib.py

+# 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, unused-argument
+"""Backend compiler related feature registration"""
+from __future__ import absolute_import
+
+import topi
+from .. import op as reg
+from ..op import OpPattern
+
+
+# adaptive_max_pool2d
+@reg.register_schedule("contrib.adaptive_max_pool2d")
+def schedule_adaptive_max_pool2d(_, outs, target):
+    """Schedule definition of adaptive_max_pool2d"""
+    with target:
+        return topi.generic.schedule_adaptive_pool(outs)
+
+reg.register_pattern("contrib.adaptive_max_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
+
+
+# adaptive_avg_pool2d
+@reg.register_schedule("contrib.adaptive_avg_pool2d")
+def schedule_adaptive_avg_pool2d(_, outs, target):
+    """Schedule definition of adaptive_avg_pool2d"""
+    with target:
+        return topi.generic.schedule_adaptive_pool(outs)
+
+reg.register_pattern("contrib.adaptive_avg_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)

python/tvm/relay/op/contrib/_make.py

+# 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.
+"""Constructor APIs"""
+from ...._ffi.function import _init_api
+
+_init_api("relay.op.contrib._make", __name__)

python/tvm/relay/op/contrib/contrib.py

+# 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, too-many-lines
+"""Contrib operations."""
+from __future__ import absolute_import as _abs
+from . import _make
+
+
+def 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.adaptive_max_pool2d(data, output_size, layout)
+
+def 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.adaptive_avg_pool2d(data, output_size, layout)

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

@@ -247,7 +247,7 @@ reg.register_pattern("nn.avg_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 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,12 @@ reg.register_pattern("nn.global_max_pool2d", OpPattern.OUT_ELEMWISE_FUSABLE)
 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)
 
+
 # leaky_relu
 reg.register_schedule("nn.leaky_relu", schedule_broadcast)
 reg.register_pattern("nn.leaky_relu", OpPattern.ELEMWISE)

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.contrib.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("contrib.adaptive_avg_pool2d");
+  return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_API("relay.op.contrib._make.adaptive_avg_pool2d")
+.set_body_typed(MakeAdaptiveAvgPool2D);
+
+RELAY_REGISTER_OP("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.contrib.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("contrib.adaptive_max_pool2d");
+  return CallNode::make(op, {data}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_API("relay.op.contrib._make.adaptive_max_pool2d")
+.set_body_typed(MakeAdaptiveMaxPool2D);
+
+RELAY_REGISTER_OP("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/frontend/mxnet/test_forward.py

@@ -170,6 +170,14 @@ def test_forward_pooling():
     mx_sym = mx.sym.Pooling(data, kernel=(3, 3), pad=(1, 1), pool_type='max')
     verify_mxnet_frontend_impl(mx_sym, (1, 20, 8, 8), (1, 20, 8, 8))
 
+def test_forward_adaptive_pooling():
+    data = mx.sym.var('data')
+    mx_sym = mx.sym.contrib.AdaptiveAvgPooling2D(data, output_size=(1,))
+    verify_mxnet_frontend_impl(mx_sym, (1, 20, 8, 8), (1, 20, 1, 1))
+
+    mx_sym = mx.sym.contrib.AdaptiveAvgPooling2D(data, output_size=(3, 3))
+    verify_mxnet_frontend_impl(mx_sym, (1, 20, 8, 8), (1, 20, 3, 3))
+
 def test_forward_lrn():
     data = mx.sym.var('data')
     mx_sym = mx.sym.LRN(data, alpha=2, beta=2, knorm=1, nsize=5)
@@ -590,6 +598,7 @@ if __name__ == '__main__':
     test_forward_split_squeeze()
     test_forward_expand_dims()
     test_forward_pooling()
+    test_forward_adaptive_pooling()
     test_forward_lrn()
     test_forward_ones()
     test_forward_zeros()

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.contrib.adaptive_avg_pool2d if pool_type == "avg" else relay.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"))

topi/include/topi/nn/pooling.h

@@ -231,6 +231,120 @@ inline Tensor pool(const Tensor& x,
                    count_include_pad);
 }
 
+
+inline Expr start_index(const Var& out_index,
+                        const Expr& odim,
+                        const Expr& idim) {
+  return out_index * idim / odim;
+}
+
+inline Expr end_index(const Var& out_index,
+                      const Expr& odim,
+                      const Expr& idim) {
+  Expr tmp = (out_index + 1) * idim / odim;
+  return tvm::ir::Select::make((out_index + 1) * idim % odim == 0,
+                               tmp, tmp + 1);
+}
+
+/*!
+* \brief Perform adaptive pooling on height and width dimension of data.
+*
+* \param x The input tensor
+* \param output_size Vector of two ints: {output_height, output_width}
+* \param pool_type The type of pooling operator
+* \param height_axis index of the height dimension
+* \param width_axis index of the width dimension
+*
+* \return The output tensor in same layout order
+*/
+inline Tensor adaptive_pool_impl(const Tensor& x,
+                                 const Array<Expr>& output_size,
+                                 PoolType pool_type,
+                                 const size_t height_axis,
+                                 const size_t width_axis) {
+  CHECK_EQ(output_size.size(), 2) << "Pooling kernel_size must have 2 elements";
+
+  auto height = x->shape[height_axis];
+  auto width = x->shape[width_axis];
+
+  auto out_height = output_size[0];
+  auto out_width = output_size[1];
+  Array<Expr> out_shape = x->shape;
+  out_shape.Set(height_axis, out_height);
+  out_shape.Set(width_axis, out_width);
+
+  if (pool_type == kMaxPool) {
+    return tvm::compute(out_shape, [&](const Array<Var>& output) {
+      Array<Expr> indices;
+      for (const Var& var : output) indices.push_back(var);
+      auto i_start_h = start_index(output[height_axis], out_height, height);
+      auto i_end_h = end_index(output[height_axis], out_height, height);
+      auto i_start_w = start_index(output[width_axis], out_width, width);
+      auto i_end_w = end_index(output[width_axis], out_width, width);
+      auto dheight = tvm::reduce_axis(Range(0, i_end_h - i_start_h), "rv1");
+      auto dwidth = tvm::reduce_axis(Range(0, i_end_w - i_start_w), "rv2");
+      indices.Set(height_axis, i_start_h + dheight);
+      indices.Set(width_axis, i_start_w + dwidth);
+      return tvm::max(x(indices), { dheight, dwidth });  // NOLINT(*)
+    }, "tensor", "adaptive_pool_max");
+  } else if (pool_type == kAvgPool) {
+    return tvm::compute(out_shape, [&](const Array<Var>& output) {
+      Array<Expr> indices;
+      for (const Var& var : output) indices.push_back(var);
+      auto i_start_h = start_index(output[height_axis], out_height, height);
+      auto i_end_h = end_index(output[height_axis], out_height, height);
+      auto i_start_w = start_index(output[width_axis], out_width, width);
+      auto i_end_w = end_index(output[width_axis], out_width, width);
+      auto divide_factor = tvm::cast(x->dtype, (i_end_h - i_start_h)
+                                               * (i_end_w - i_start_w));
+      auto dheight = tvm::reduce_axis(Range(0, i_end_h - i_start_h), "rv1");
+      auto dwidth = tvm::reduce_axis(Range(0, i_end_w - i_start_w), "rv2");
+      indices.Set(height_axis, i_start_h + dheight);
+      indices.Set(width_axis, i_start_w + dwidth);
+      return tvm::sum(x(indices) / divide_factor, { dheight, dwidth });
+    }, "tensor", "adaptive_pool_avg");
+  } else {
+    LOG(ERROR) << "Unrecognized pool_type: " << pool_type;
+    return x;
+  }
+}
+
+/*!
+* \brief Adaptively perform pooling on height and width dimension of data.
+*        The pooling kernel and stride sizes are automatically chosen for desired output sizes.
+*        It decides the height and width dimension according to the layout string,
+*        in which 'W' and 'H' means width and height respectively.
+*        Width and height dimension cannot be split.
+*        For example, NCHW, NCHW16c, etc. are valid for pool,
+*        while NCHW16w, NCHW16h are not.
+*        See \a layout for more information of the layout string convention.
+*
+* \param x The input tensor
+* \param output_size Vector of two ints: {output_height, output_width}
+* \param pool_type The type of pooling operator
+* \param layout The input layout. Pooling supports any layout as long as 'H' and 'W' appear.
+*        The layout is supposed to be composed of upper cases, lower cases and (optional) numbers,
+*        where upper case indicates a dimension and
+*        the corresponding lower case (with factor size) indicates the split dimension.
+*        For example, NCHW16c can describe a 5-D tensor of
+*        [batch_size, channel, height, width, channel_block].
+*        (in which factor size `16` will not be used in pooling but for other operators,
+*        it can be used to decide the output shape).
+*        Since pooling does not care about the factor size of dimensions
+*        other than `H` and `W`, one can pass `NCHWc` as well.
+*
+* \return The output tensor in same layout order
+*/
+inline Tensor adaptive_pool(const Tensor& x,
+                            const Array<Expr>& output_size,
+                            PoolType pool_type,
+                            const std::string& layout = "NCHW") {
+  int height_axis = -1, width_axis = -1;
+  CHECK(find_height_width(layout, &height_axis, &width_axis))
+    << "Unsupported layout " << layout;
+  return adaptive_pool_impl(x, output_size, pool_type, height_axis, width_axis);
+}
+
 /*!
 * \brief Perform global pooling on height and width dimension of data.
 *        It decides the height and width dimension according to the layout string,
@@ -259,49 +373,7 @@ inline Tensor pool(const Tensor& x,
 inline Tensor global_pool(const Tensor& x,
                           PoolType pool_type,
                           const std::string& layout = "NCHW") {
-  CHECK(x->shape.size() >= 2) << "Pooling input must >= 2-D (H, W)";
-
-  int height_axis = -1, width_axis = -1;
-  CHECK(find_height_width(layout, &height_axis, &width_axis))
-    << "Unsupported layout " << layout;
-
-  Array<Expr> out_shape = x->shape;
-  out_shape.Set(height_axis, 1);
-  out_shape.Set(width_axis, 1);
-
-  auto height = x->shape[height_axis];
-  auto width = x->shape[width_axis];
-
-  auto dheight = tvm::reduce_axis(Range(0, height), "rv1");
-  auto dwidth = tvm::reduce_axis(Range(0, width), "rv2");
-
-  if (pool_type == kMaxPool) {
-    return tvm::compute(out_shape,
-      [&](const Array<Var>& output) {
-        Array<Expr> indices;
-        for (const Var& var : output) indices.push_back(var);
-        indices.Set(height_axis, dheight);
-        indices.Set(width_axis, dwidth);
-        return tvm::max(x(indices), { dheight, dwidth });  // NOLINT(*)
-      }, "tensor", "global_pool_max");
-  } else if (pool_type == kAvgPool) {
-    auto tsum = tvm::compute(out_shape,
-      [&](const Array<Var>& output) {
-        Array<Expr> indices;
-        for (const Var& var : output) indices.push_back(var);
-        indices.Set(height_axis, dheight);
-        indices.Set(width_axis, dwidth);
-        return tvm::sum(x(indices), { dheight, dwidth });
-      }, "tensor", "global_pool_sum");
-
-    return tvm::compute(out_shape,
-      [&](const Array<Var>& output) {
-        return tsum(output) / tvm::cast(x->dtype, height * width);
-      }, "tensor", kElementWise);
-  } else {
-    LOG(ERROR) << "Unrecognized pool_type: " << pool_type;
-    return x;
-  }
+  return adaptive_pool(x, Array<Expr>{1, 1}, pool_type, layout);
 }
 
 }  // namespace nn

topi/python/topi/cuda/__init__.py

@@ -12,7 +12,7 @@ from .reduction import schedule_reduce
 from .softmax import schedule_softmax
 from .injective import schedule_injective, schedule_elemwise, schedule_broadcast
 from .dense import schedule_dense
-from .pooling import schedule_pool, schedule_global_pool
+from .pooling import schedule_pool, schedule_adaptive_pool
 from .extern import schedule_extern
 from .nn import schedule_lrn, schedule_l2_normalize
 from .batch_matmul import schedule_batch_matmul

topi/python/topi/cuda/pooling.py

@@ -20,23 +20,26 @@ import tvm
 from .. import tag
 from .. import generic
 
-@generic.schedule_global_pool.register(["cuda", "gpu"])
-def schedule_global_pool(outs):
-    """Schedule for global_pool.
+
+
+@generic.schedule_adaptive_pool.register(["cuda", "gpu"])
+def schedule_adaptive_pool(outs):
+    """Schedule for adaptive_pool.
 
     Parameters
     ----------
     outs: Array of Tensor
-        The computation graph description of global_pool
+        The computation graph description of adaptive_pool
         in the format of an array of tensors.
 
     Returns
     -------
     s: Schedule
-        The computation schedule for global_pool.
+        The computation schedule for adaptive_pool.
     """
     outs = [outs] if isinstance(outs, tvm.tensor.Tensor) else outs
     s = tvm.create_schedule([x.op for x in outs])
+
     def _schedule(Pool):
         num_thread = 8
         block_x = tvm.thread_axis("blockIdx.x")
@@ -73,7 +76,7 @@ def schedule_global_pool(outs):
                 if tensor.op.input_tensors and tensor.op not in scheduled_ops:
                     traverse(tensor.op)
         # schedule global_pool
-        elif OP.tag.startswith('global_pool'):
+        elif OP.tag.startswith('adaptive_pool'):
             Pool = OP.output(0)
             _schedule(Pool)
         else:

topi/python/topi/generic/nn.py

@@ -403,14 +403,14 @@ def schedule_pool(outs, layout):
     return _default_schedule(outs, False)
 
 
-@tvm.target.override_native_generic_func("schedule_global_pool")
-def schedule_global_pool(outs):
-    """Schedule for global pool
+@tvm.target.override_native_generic_func("schedule_adaptive_pool")
+def schedule_adaptive_pool(outs):
+    """Schedule for adaptive pool
 
     Parameters
     ----------
     outs: Array of Tensor
-          The computation graph description of global pool
+          The computation graph description of adaptive pool
           in the format of an array of tensors.
 
     Returns
@@ -420,6 +420,7 @@ def schedule_global_pool(outs):
     """
     return _default_schedule(outs, False)
 
+
 @tvm.target.override_native_generic_func("schedule_binarize_pack")
 def schedule_binarize_pack(outs):
     """Schedule for binarize_pack

topi/python/topi/hls/nn.py

@@ -360,14 +360,14 @@ def schedule_pool(outs, layout):
     return s
 
 
-@generic.schedule_global_pool.register(["hls"])
-def schedule_global_pool(outs):
-    """Schedule for global pool
+@generic.schedule_adaptive_pool.register(["hls"])
+def schedule_adaptive_pool(outs):
+    """Schedule for adaptive_pool
 
     Parameters
     ----------
     outs: Array of Tensor
-          The computation graph description of global pool
+          The computation graph description of adaptive_pool
           in the format of an array of tensors.
 
     Returns
@@ -389,7 +389,7 @@ def schedule_global_pool(outs):
                 if tensor.op.input_tensors:
                     traverse(tensor.op)
         # schedule global_pool
-        elif OP.tag.startswith('global_pool'):
+        elif OP.tag.startswith('adaptive_pool'):
             Pool = OP.output(0)
             if not Pool.op in s.outputs:
                 Out = outs[0].op.output(0)

topi/python/topi/nn/pooling.py

@@ -113,3 +113,44 @@ def pool(data,
     """
     return cpp.nn.pool(data, kernel, stride, padding,
                        POOL_TYPE_CODE[pool_type], ceil_mode, layout, count_include_pad)
+
+def adaptive_pool(data,
+                  output_size,
+                  pool_type,
+                  layout="NCHW"):
+    """Perform pooling on height and width dimension of data.
+       The pooling kernel and stride sizes are automatically chosen for desired
+       output sizes.
+       It decides the height and width dimension according to the layout string,
+       in which 'W' and 'H' means width and height respectively.
+       Width and height dimension cannot be split.
+       For example, NCHW, NCHW16c, etc. are valid for pool,
+       while NCHW16w, NCHW16h are not.
+       See parameter `layout` for more information of the layout string convention.
+
+    Parameters
+    ----------
+    data : tvm.Tensor
+        n-D with shape of layout
+
+    output_size : tuple of int
+        output height and width.
+
+    pool_type : str
+        Pool type, 'max' or 'avg'
+
+    layout: string
+        Layout of the input data.
+        The layout is supposed to be composed of upper cases, lower cases and numbers,
+        where upper case indicates a dimension and
+        the corresponding lower case with factor size indicates the split dimension.
+        For example, NCHW16c can describe a 5-D tensor of
+        [batch_size, channel, height, width, channel_block],
+        in which channel_block=16 is a split of dimension channel.
+
+    Returns
+    -------
+    output : tvm.Tensor
+        n-D in the same layout
+    """
+    return cpp.nn.adaptive_pool(data, output_size, POOL_TYPE_CODE[pool_type], layout)

topi/python/topi/opengl/__init__.py

@@ -6,4 +6,4 @@ from .conv2d_nchw import schedule_conv2d_nchw
 from .injective import schedule_injective, schedule_elemwise, schedule_broadcast
 from .softmax import schedule_softmax
 from .dense import schedule_dense
-from .pooling import schedule_pool, schedule_global_pool
+from .pooling import schedule_pool, schedule_adaptive_pool

topi/python/topi/opengl/pooling.py

@@ -20,9 +20,9 @@ import tvm
 from .. import tag
 from .. import generic
 
-@generic.schedule_global_pool.register(["opengl"])
-def schedule_global_pool(outs):
-    """Schedule for global_pool.
+@generic.schedule_adaptive_pool.register(["opengl"])
+def schedule_adaptive_pool(outs):
+    """Schedule for adaptive pool.
 
     Parameters
     ----------
@@ -33,7 +33,7 @@ def schedule_global_pool(outs):
     Returns
     -------
     s: Schedule
-        The computation schedule for global_pool.
+        The computation schedule for adaptive pool.
     """
     outs = [outs] if isinstance(outs, tvm.tensor.Tensor) else outs
     s = tvm.create_schedule([x.op for x in outs])
@@ -57,7 +57,7 @@ def schedule_global_pool(outs):
                 if tensor.op.input_tensors and tensor.op not in scheduled_ops:
                     traverse(tensor.op)
         # schedule global_pool
-        elif OP.tag.startswith('global_pool'):
+        elif OP.tag.startswith('adaptive_pool'):
             Pool = OP.output(0)
             _schedule(Pool)
         else:

topi/python/topi/x86/__init__.py

@@ -7,7 +7,7 @@ from .binarize_pack import schedule_binarize_pack
 from .binary_dense import schedule_binary_dense
 from .nn import *
 from .injective import *
-from .pooling import schedule_pool, schedule_global_pool
+from .pooling import schedule_pool, schedule_adaptive_pool
 from .bitserial_conv2d import schedule_bitserial_conv2d
 from .bitserial_dense import schedule_bitserial_dense
 from .depthwise_conv2d import schedule_depthwise_conv2d_NCHWc

topi/python/topi/x86/pooling.py

@@ -110,14 +110,14 @@ def schedule_pool(outs, layout):
     return s
 
 
-@generic.schedule_global_pool.register(["cpu"])
-def schedule_global_pool(outs):
-    """Schedule for global pool
+@generic.schedule_adaptive_pool.register(["cpu"])
+def schedule_adaptive_pool(outs):
+    """Schedule for adaptive pool
 
     Parameters
     ----------
     outs: Array of Tensor
-          The computation graph description of pool
+          The computation graph description of adaptive pool
           in the format of an array of tensors.
 
     Returns
@@ -139,7 +139,7 @@ def schedule_global_pool(outs):
                 if tensor.op.input_tensors and tensor.op not in scheduled_ops:
                     traverse(tensor.op)
         # schedule pool
-        elif OP.tag.startswith('global_pool'):
+        elif OP.tag.startswith('adaptive_pool'):
             Pool = OP.output(0)
             _parallel_sch(s[Pool], outs[0].shape)
         else:

topi/src/topi.cc

@@ -455,6 +455,13 @@ TVM_REGISTER_GLOBAL("topi.nn.global_pool")
                         static_cast<nn::PoolType>(static_cast<int>(args[1])));
   });
 
+TVM_REGISTER_GLOBAL("topi.nn.adaptive_pool")
+.set_body([](TVMArgs args, TVMRetValue *rv) {
+  *rv = nn::adaptive_pool(args[0], args[1],
+                          static_cast<nn::PoolType>(static_cast<int>(args[2])),
+                          args[3]);
+});
+
 /* Ops from nn/softmax.h */
 TVM_REGISTER_GLOBAL("topi.nn.softmax")
 .set_body([](TVMArgs args, TVMRetValue *rv) {

topi/tests/python/test_topi_pooling.py

@@ -120,7 +120,7 @@ def verify_global_pool(n, c, h, w, pool_type):
             return
         print("Running on target: %s" % device)
         with tvm.target.create(device):
-            s = topi.generic.schedule_global_pool(B)
+            s = topi.generic.schedule_adaptive_pool(B)
         a = tvm.nd.array(a_np, ctx)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), ctx)
         f = tvm.build(s, [A, B], device)
@@ -136,7 +136,58 @@ def test_global_pool():
     verify_global_pool(1, 1024, 7, 7, 'max')
     verify_global_pool(4, 1024, 7, 7, 'max')
 
+def verify_adaptive_pool(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])
+
+    data = tvm.placeholder(dshape, name="data", dtype=dtype)
+    out = topi.nn.adaptive_pool(data, out_size, pool_type, layout)
+    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_adaptive_pool(out)
+        a = tvm.nd.array(np_data, ctx)
+        b = tvm.nd.array(np.zeros(get_const_tuple(oshape), dtype=out.dtype), ctx)
+        f = tvm.build(s, [data, out], device)
+        f(a, b)
+        tvm.testing.assert_allclose(b.asnumpy(), np_out, rtol=1e-5)
+
+    for device in get_all_backend():
+        check_device(device)
+
+def test_adaptive_pool():
+    verify_adaptive_pool((1, 3, 224, 224), (1, 1), "max")
+    verify_adaptive_pool((1, 3, 224, 224), (1, 1), "avg")
+    verify_adaptive_pool((1, 14, 56, 78), (34, 13), "max")
+    verify_adaptive_pool((1, 5, 46, 97), (4, 96), "avg")
+
 
 if __name__ == "__main__":
     test_pool()
     test_global_pool()
+    test_adaptive_pool()