From 147ea3b0ca147b527086228d524a2f68f872112d Mon Sep 17 00:00:00 2001 From: Yao Wang Date: Wed, 8 May 2019 17:21:41 -0700 Subject: [PATCH] [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 | 1 + docs/langref/relay_op.rst | 4 + include/tvm/relay/attrs/nn.h | 16 +++ nnvm/python/nnvm/top/nn.py | 4 +- python/tvm/relay/__init__.py | 1 + python/tvm/relay/contrib.py | 20 ++++ python/tvm/relay/frontend/mxnet.py | 5 +- python/tvm/relay/op/__init__.py | 1 + python/tvm/relay/op/contrib/__init__.py | 21 ++++ python/tvm/relay/op/contrib/_contrib.py | 43 +++++++ python/tvm/relay/op/contrib/_make.py | 20 ++++ python/tvm/relay/op/contrib/contrib.py | 113 +++++++++++++++++++ python/tvm/relay/op/nn/_nn.py | 5 +- src/relay/op/nn/pooling.cc | 167 +++++++++++++++++++++++++++- tests/python/frontend/mxnet/test_forward.py | 9 ++ tests/python/relay/test_op_level10.py | 43 +++++++ tests/python/relay/test_op_level2.py | 1 - topi/include/topi/nn/pooling.h | 158 +++++++++++++++++++------- topi/python/topi/cuda/__init__.py | 2 +- topi/python/topi/cuda/pooling.py | 15 ++- topi/python/topi/generic/nn.py | 9 +- topi/python/topi/hls/nn.py | 10 +- topi/python/topi/nn/pooling.py | 41 +++++++ topi/python/topi/opengl/__init__.py | 2 +- topi/python/topi/opengl/pooling.py | 10 +- topi/python/topi/x86/__init__.py | 2 +- topi/python/topi/x86/pooling.py | 10 +- topi/src/topi.cc | 7 ++ topi/tests/python/test_topi_pooling.py | 53 ++++++++- 29 files changed, 710 insertions(+), 83 deletions(-) create mode 100644 python/tvm/relay/contrib.py create mode 100644 python/tvm/relay/op/contrib/__init__.py create mode 100644 python/tvm/relay/op/contrib/_contrib.py create mode 100644 python/tvm/relay/op/contrib/_make.py create mode 100644 python/tvm/relay/op/contrib/contrib.py diff --git a/docs/api/python/topi.rst b/docs/api/python/topi.rst index 222b334..eaa5dac 100644 --- a/docs/api/python/topi.rst +++ b/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 diff --git a/docs/langref/relay_op.rst b/docs/langref/relay_op.rst index 4719aba..cd567729 100644 --- a/docs/langref/relay_op.rst +++ b/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 diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index 2049a8f..8a1aca0 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -332,6 +332,22 @@ struct GlobalPool2DAttrs : public tvm::AttrsNode { } }; +/*! \brief Attributes for adaptive pool operator */ +struct AdaptivePool2DAttrs : public tvm::AttrsNode { + Array output_size; + std::string layout; + + TVM_DECLARE_ATTRS(AdaptivePool2DAttrs, "relay.attrs.AdaptivePool2DAttrs") { + TVM_ATTR_FIELD(output_size).set_default(Array({})) + .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 { diff --git a/nnvm/python/nnvm/top/nn.py b/nnvm/python/nnvm/top/nn.py index 71367b4..13964f4 100644 --- a/nnvm/python/nnvm/top/nn.py +++ b/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) diff --git a/python/tvm/relay/__init__.py b/python/tvm/relay/__init__.py index 80555d3..6201681 100644 --- a/python/tvm/relay/__init__.py +++ b/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 diff --git a/python/tvm/relay/contrib.py b/python/tvm/relay/contrib.py new file mode 100644 index 0000000..d22c676 --- /dev/null +++ b/python/tvm/relay/contrib.py @@ -0,0 +1,20 @@ +# 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 * diff --git a/python/tvm/relay/frontend/mxnet.py b/python/tvm/relay/frontend/mxnet.py index b93bd5b..1a4d52f 100644 --- a/python/tvm/relay/frontend/mxnet.py +++ b/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): diff --git a/python/tvm/relay/op/__init__.py b/python/tvm/relay/op/__init__.py index 3bea795..a27ab1d 100644 --- a/python/tvm/relay/op/__init__.py +++ b/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 diff --git a/python/tvm/relay/op/contrib/__init__.py b/python/tvm/relay/op/contrib/__init__.py new file mode 100644 index 0000000..3159006 --- /dev/null +++ b/python/tvm/relay/op/contrib/__init__.py @@ -0,0 +1,21 @@ +# 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 diff --git a/python/tvm/relay/op/contrib/_contrib.py b/python/tvm/relay/op/contrib/_contrib.py new file mode 100644 index 0000000..f0df756 --- /dev/null +++ b/python/tvm/relay/op/contrib/_contrib.py @@ -0,0 +1,43 @@ +# 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) diff --git a/python/tvm/relay/op/contrib/_make.py b/python/tvm/relay/op/contrib/_make.py new file mode 100644 index 0000000..42d7175 --- /dev/null +++ b/python/tvm/relay/op/contrib/_make.py @@ -0,0 +1,20 @@ +# 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__) diff --git a/python/tvm/relay/op/contrib/contrib.py b/python/tvm/relay/op/contrib/contrib.py new file mode 100644 index 0000000..1f073d4 --- /dev/null +++ b/python/tvm/relay/op/contrib/contrib.py @@ -0,0 +1,113 @@ +# 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) diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index 5e9d5d7..6c8f8f8 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/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) diff --git a/src/relay/op/nn/pooling.cc b/src/relay/op/nn/pooling.cc index df238b3..4dd763b 100644 --- a/src/relay/op/nn/pooling.cc +++ b/src/relay/op/nn/pooling.cc @@ -72,7 +72,6 @@ bool Pool2DRel(const Array& 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(); @@ -284,7 +283,6 @@ bool GlobalPool2DRel(const Array& types, const auto* data = types[0].as(); 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(); @@ -393,5 +391,170 @@ RELAY_REGISTER_OP("nn.global_max_pool2d") Pool2DInferCorrectLayout) .set_attr("FTVMCompute", GlobalPool2DCompute); + +// relay.nn.adaptive_pool_2d +TVM_REGISTER_NODE_TYPE(AdaptivePool2DAttrs); + +bool AdaptivePool2DRel(const Array& types, + int num_inputs, + const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 2); + const auto* data = types[0].as(); + 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(); + 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 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 +Array AdaptivePool2DCompute(const Attrs& attrs, + const Array& inputs, + const Type& out_type, + const Target& target) { + static const Layout kNCHW("NCHW"); + const auto* param = attrs.as(); + 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{ + topi::nn::adaptive_pool(inputs[0], Array{ output_height, output_width }, + mode, layout.name()) }; +} + +// relay.contrib.adaptive_avg_pool2d +Expr MakeAdaptiveAvgPool2D(Expr data, + Array output_size, + std::string layout) { + auto attrs = make_node(); + 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", + Pool2DInferCorrectLayout) +.set_attr("FTVMCompute", AdaptivePool2DCompute); + + +// relay.contrib.adaptive_max_pool2d +Expr MakeAdaptiveMaxPool2D(Expr data, + Array output_size, + std::string layout) { + auto attrs = make_node(); + 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", + Pool2DInferCorrectLayout) +.set_attr("FTVMCompute", AdaptivePool2DCompute); + } // namespace relay } // namespace tvm diff --git a/tests/python/frontend/mxnet/test_forward.py b/tests/python/frontend/mxnet/test_forward.py index 067c356..e75e60d 100644 --- a/tests/python/frontend/mxnet/test_forward.py +++ b/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() diff --git a/tests/python/relay/test_op_level10.py b/tests/python/relay/test_op_level10.py index 667dd87..244744c 100644 --- a/tests/python/relay/test_op_level10.py +++ b/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() diff --git a/tests/python/relay/test_op_level2.py b/tests/python/relay/test_op_level2.py index 88963a6..a535045 100644 --- a/tests/python/relay/test_op_level2.py +++ b/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")) diff --git a/topi/include/topi/nn/pooling.h b/topi/include/topi/nn/pooling.h index 8648c01..08cc935 100644 --- a/topi/include/topi/nn/pooling.h +++ b/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& 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 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& output) { + Array 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& output) { + Array 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& 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 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& output) { - Array 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& output) { - Array 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& 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{1, 1}, pool_type, layout); } } // namespace nn diff --git a/topi/python/topi/cuda/__init__.py b/topi/python/topi/cuda/__init__.py index 65ed0ff..526429b 100644 --- a/topi/python/topi/cuda/__init__.py +++ b/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 diff --git a/topi/python/topi/cuda/pooling.py b/topi/python/topi/cuda/pooling.py index ac3644d..2d12c4a 100644 --- a/topi/python/topi/cuda/pooling.py +++ b/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: diff --git a/topi/python/topi/generic/nn.py b/topi/python/topi/generic/nn.py index db77f37..db1c772 100644 --- a/topi/python/topi/generic/nn.py +++ b/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 diff --git a/topi/python/topi/hls/nn.py b/topi/python/topi/hls/nn.py index 5ce8394..adb8862 100644 --- a/topi/python/topi/hls/nn.py +++ b/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) diff --git a/topi/python/topi/nn/pooling.py b/topi/python/topi/nn/pooling.py index f7f59ca..fb980aa 100644 --- a/topi/python/topi/nn/pooling.py +++ b/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) diff --git a/topi/python/topi/opengl/__init__.py b/topi/python/topi/opengl/__init__.py index c8f20b9..37eac44 100644 --- a/topi/python/topi/opengl/__init__.py +++ b/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 diff --git a/topi/python/topi/opengl/pooling.py b/topi/python/topi/opengl/pooling.py index 99c4dec..56f0b08 100644 --- a/topi/python/topi/opengl/pooling.py +++ b/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: diff --git a/topi/python/topi/x86/__init__.py b/topi/python/topi/x86/__init__.py index a414e3f..cce816d 100644 --- a/topi/python/topi/x86/__init__.py +++ b/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 diff --git a/topi/python/topi/x86/pooling.py b/topi/python/topi/x86/pooling.py index 4cf213a..816e03c 100644 --- a/topi/python/topi/x86/pooling.py +++ b/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: diff --git a/topi/src/topi.cc b/topi/src/topi.cc index d486e7b..1585d87 100644 --- a/topi/src/topi.cc +++ b/topi/src/topi.cc @@ -455,6 +455,13 @@ TVM_REGISTER_GLOBAL("topi.nn.global_pool") static_cast(static_cast(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(static_cast(args[2])), + args[3]); +}); + /* Ops from nn/softmax.h */ TVM_REGISTER_GLOBAL("topi.nn.softmax") .set_body([](TVMArgs args, TVMRetValue *rv) { diff --git a/topi/tests/python/test_topi_pooling.py b/topi/tests/python/test_topi_pooling.py index 36dcd0f..bba14ec 100644 --- a/topi/tests/python/test_topi_pooling.py +++ b/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() -- 2.7.4