topi.nn.dilate
topi.nn.pool
topi.nn.global_pool
+ topi.nn.adaptive_pool
topi.nn.upsampling
topi.nn.softmax
topi.nn.dense
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
.. 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
}
};
+/*! \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> {
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)
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)
from . import nn
from . import annotation
from . import vision
+from . import contrib
from . import image
from . import frontend
from . import backend
--- /dev/null
+# 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 *
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):
from . import annotation
from . import image
from . import vision
+from . import contrib
from . import op_attrs
--- /dev/null
+# 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
--- /dev/null
+# 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)
--- /dev/null
+# 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__)
--- /dev/null
+# 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)
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)
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)
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>();
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>();
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
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)
test_forward_split_squeeze()
test_forward_expand_dims()
test_forward_pooling()
+ test_forward_adaptive_pooling()
test_forward_lrn()
test_forward_ones()
test_forward_zeros()
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()
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"))
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,
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
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
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")
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:
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
"""
return _default_schedule(outs, False)
+
@tvm.target.override_native_generic_func("schedule_binarize_pack")
def schedule_binarize_pack(outs):
"""Schedule for binarize_pack
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
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)
"""
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)
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
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
----------
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])
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:
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
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
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:
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) {
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)
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()
projects / platform / upstream / tvm.git / commitdiff