chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,406 @@
|
||||
# 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.
|
||||
"""TVM operator pooling compute."""
|
||||
|
||||
from .. import cpp
|
||||
|
||||
POOL_TYPE_CODE = {"avg": 0, "max": 1}
|
||||
|
||||
|
||||
def global_pool(data, pool_type, layout="NCHW"):
|
||||
"""Perform global pooling on height and width dimension of data.
|
||||
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.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
layout : str
|
||||
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.te.Tensor
|
||||
n-D in same layout with height and width dimension size of 1.
|
||||
e.g., for NCHW, the output shape will be [batch, channel, 1, 1]
|
||||
"""
|
||||
return cpp.nn.global_pool(data, POOL_TYPE_CODE[pool_type], layout)
|
||||
|
||||
|
||||
def pool_grad(
|
||||
grads,
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
padding,
|
||||
pool_type,
|
||||
ceil_mode=False,
|
||||
count_include_pad=True,
|
||||
layout="NCHW",
|
||||
):
|
||||
"""Gradient of pooling on height and width dimension of data.
|
||||
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
|
||||
----------
|
||||
grads : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
data : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
kernel : list/tuple of two ints
|
||||
Kernel size, [kernel_height, kernel_width]
|
||||
|
||||
stride : list/tuple of two ints
|
||||
Stride size, [stride_height, stride_width]
|
||||
|
||||
padding : list/tuple of four ints
|
||||
Pad size, [pad_top, pad_left, pad_bottom, pad_right]]
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
ceil_mode : bool
|
||||
Whether to use ceil when calculating output size.
|
||||
|
||||
count_include_pad: bool
|
||||
Whether include padding in the calculation when pool_type is '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.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
return cpp.nn.pool_grad(
|
||||
grads,
|
||||
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.te.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.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
return cpp.nn.adaptive_pool(data, output_size, POOL_TYPE_CODE[pool_type], layout)
|
||||
|
||||
|
||||
def adaptive_pool1d(data, output_size, pool_type, layout="NCW"):
|
||||
"""Perform pooling on three dimensional data.
|
||||
See the two dimensional version above for details.
|
||||
"""
|
||||
return cpp.nn.adaptive_pool1d(data, output_size, POOL_TYPE_CODE[pool_type], layout)
|
||||
|
||||
|
||||
def adaptive_pool3d(data, output_size, pool_type, layout="NCDHW"):
|
||||
"""Perform pooling on three dimensional data.
|
||||
See the two dimensional version above for details.
|
||||
"""
|
||||
return cpp.nn.adaptive_pool3d(data, output_size, POOL_TYPE_CODE[pool_type], layout)
|
||||
|
||||
|
||||
def pool1d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
pool_type,
|
||||
ceil_mode=False,
|
||||
layout="NCW",
|
||||
count_include_pad=True,
|
||||
):
|
||||
"""Perform pooling on width dimension of data.
|
||||
Width axis is determined according to the layout string.
|
||||
in which 'w' means width.
|
||||
Width dimension cannot be split.
|
||||
For example, NCW, NCW16c, etc. are valid for pool,
|
||||
while NCW16w is not.
|
||||
See parameter `layout` for more information of the layout string convention.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
kernel : list/tuple of one int or int
|
||||
Kernel size, [kernel_width]
|
||||
|
||||
stride : list/tuple of one int or int
|
||||
Stride size, [stride_width]
|
||||
|
||||
dilation: list/tuple of two ints
|
||||
Dilation size, [dilation_height, dilation_width]
|
||||
|
||||
padding : list/tuple of two ints
|
||||
Pad size, [pad_left, pad_right]
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
ceil_mode : bool
|
||||
Whether to use ceil when calculating output size.
|
||||
|
||||
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, NCW16c can describe a 4-D tensor of
|
||||
[batch_size, channel, width, channel_block],
|
||||
in which channel_block=16 is a split of dimension channel.
|
||||
|
||||
count_include_pad: bool
|
||||
Whether include padding in the calculation when pool_type is 'avg'
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
if isinstance(kernel, int):
|
||||
kernel = [
|
||||
kernel,
|
||||
]
|
||||
if isinstance(stride, int):
|
||||
stride = [
|
||||
stride,
|
||||
]
|
||||
return cpp.nn.pool1d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
POOL_TYPE_CODE[pool_type],
|
||||
ceil_mode,
|
||||
layout,
|
||||
count_include_pad,
|
||||
)
|
||||
|
||||
|
||||
def pool2d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
pool_type,
|
||||
ceil_mode=False,
|
||||
layout="NCHW",
|
||||
count_include_pad=True,
|
||||
):
|
||||
"""Perform pooling on height and width dimension of data.
|
||||
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.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
kernel : list/tuple of two ints
|
||||
Kernel size, [kernel_height, kernel_width]
|
||||
|
||||
stride : list/tuple of two ints
|
||||
Stride size, [stride_height, stride_width]
|
||||
|
||||
dilation: list/tuple of two ints
|
||||
Dilation size, [dilation_height, dilation_width]
|
||||
|
||||
padding : list/tuple of four ints
|
||||
Pad size, [pad_top, pad_left, pad_bottom, pad_right]]
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
ceil_mode : bool
|
||||
Whether to use ceil when calculating output size.
|
||||
|
||||
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.
|
||||
|
||||
count_include_pad: bool
|
||||
Whether include padding in the calculation when pool_type is 'avg'
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
return cpp.nn.pool2d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
POOL_TYPE_CODE[pool_type],
|
||||
ceil_mode,
|
||||
layout,
|
||||
count_include_pad,
|
||||
)
|
||||
|
||||
|
||||
def pool3d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
pool_type,
|
||||
ceil_mode=False,
|
||||
layout="NCDHW",
|
||||
count_include_pad=True,
|
||||
):
|
||||
"""Perform pooling on depth, height and width dimension of data.
|
||||
It decides the depth, height and width dimension according to the layout string,
|
||||
in which 'D', 'W' and 'H' means depth, width and height respectively.
|
||||
Depth, width and height dimension cannot be split.
|
||||
For example, NCDHW, NCDHW16c, etc. are valid for pool,
|
||||
while NCDHW16d, NCDHW16w, NCDHW16h are not.
|
||||
See parameter `layout` for more information of the layout string convention.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
n-D with shape of layout
|
||||
|
||||
kernel : list/tuple of three ints
|
||||
Kernel size, [kernel_depth, kernel_height, kernel_width]
|
||||
|
||||
stride : list/tuple of three ints
|
||||
Stride size, [stride_depth, stride_height, stride_width]
|
||||
|
||||
dilation: list/tuple of two ints
|
||||
Dilation size, [dilation_height, dilation_width]
|
||||
|
||||
padding : list/tuple of six ints
|
||||
Pad size, [pad_front, pad_top, pad_left, pad_back, pad_bottom, pad_right]
|
||||
|
||||
pool_type : str
|
||||
Pool type, 'max' or 'avg'
|
||||
|
||||
ceil_mode : bool
|
||||
Whether to use ceil when calculating output size.
|
||||
|
||||
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, NCDHW16c can describe a 6-D tensor of
|
||||
[batch_size, channel, depth, height, width, channel_block],
|
||||
in which channel_block=16 is a split of dimension channel.
|
||||
|
||||
count_include_pad: bool
|
||||
Whether include padding in the calculation when pool_type is 'avg'
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
n-D in the same layout
|
||||
"""
|
||||
return cpp.nn.pool3d(
|
||||
data,
|
||||
kernel,
|
||||
stride,
|
||||
dilation,
|
||||
padding,
|
||||
POOL_TYPE_CODE[pool_type],
|
||||
ceil_mode,
|
||||
layout,
|
||||
count_include_pad,
|
||||
)
|
||||
Reference in New Issue
Block a user