chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,217 @@
|
||||
# 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-variable, unused-argument
|
||||
"""Transposed 1D convolution operators (sometimes called Deconvolution)."""
|
||||
|
||||
from tvm import te
|
||||
|
||||
from ..utils import simplify
|
||||
from .dilate import dilate
|
||||
from .pad import pad
|
||||
from .utils import get_pad_tuple1d
|
||||
|
||||
|
||||
def _conv1d_transpose_ncw_preprocess(data, kernel, stride, padding, out_dtype, output_padding):
|
||||
"""Preprocess data and kernel to make the compute pattern
|
||||
of conv1d_transpose the same as conv1d.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
3-D with shape [in_channel, num_filter, filter_width]
|
||||
|
||||
stride : ints
|
||||
The spatial stride along width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : ints
|
||||
Used to recover the actual output shape in case there are more
|
||||
than one possible shape. Must be smaller than stride.
|
||||
|
||||
Returns
|
||||
-------
|
||||
data_pad : tvm.te.Tensor
|
||||
Padded input data. 3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
kernel: tvm.te.Tensor
|
||||
Transformed kernel. 3-D with shape [num_filter, in_channel, filter_width]
|
||||
"""
|
||||
# some pre-processing and prelimnary checks
|
||||
if out_dtype is None:
|
||||
out_dtype = data.dtype
|
||||
|
||||
# dilate and pad
|
||||
if isinstance(stride, tuple | list):
|
||||
stride = stride[0]
|
||||
if isinstance(output_padding, tuple | list):
|
||||
output_padding = output_padding[0]
|
||||
|
||||
_, channels_in, _ = data.shape
|
||||
_, channels_out, kernel_width = kernel.shape
|
||||
assert output_padding < stride
|
||||
channels_out = simplify(channels_out)
|
||||
data_dilate = dilate(data, [1, 1, stride], name="data_dilate")
|
||||
pad_left, pad_right = get_pad_tuple1d(padding, (kernel_width,))
|
||||
pad_left = kernel_width - 1 - pad_left
|
||||
pad_right = kernel_width - 1 - pad_right + output_padding
|
||||
data_pad = pad(data_dilate, [0, 0, pad_left], [0, 0, pad_right], name="data_pad")
|
||||
|
||||
# transform kernel layout from IOW to OIW, and rotate kernel by 180 degrees
|
||||
kernel = te.compute(
|
||||
(channels_out, channels_in, kernel_width),
|
||||
lambda o, i, w: kernel[i][o][kernel_width - 1 - w],
|
||||
name="kernel",
|
||||
)
|
||||
return data_pad, kernel
|
||||
|
||||
|
||||
def conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding):
|
||||
"""Transposed 1D convolution ncw forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
3-D with shape [in_channel, num_filter, filter_width]
|
||||
|
||||
stride : ints
|
||||
The spatial stride along width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : ints
|
||||
Used to recover the actual output shape in case there are more
|
||||
than one possible shape. Must be smaller than stride.
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
3-D with shape [batch, out_channel, out_width]
|
||||
|
||||
"""
|
||||
|
||||
batch, channels_in, _ = data.shape
|
||||
_, channels_out, kernel_width = kernel.shape
|
||||
|
||||
data_pad, transformed_kernel = _conv1d_transpose_ncw_preprocess(
|
||||
data, kernel, stride, padding, out_dtype, output_padding
|
||||
)
|
||||
|
||||
# convolution
|
||||
_, _, data_width = data_pad.shape
|
||||
out_w = simplify(data_width - kernel_width + 1)
|
||||
dc = te.reduce_axis((0, channels_in), name="dc")
|
||||
dw = te.reduce_axis((0, kernel_width), name="dw")
|
||||
output = te.compute(
|
||||
(batch, channels_out, out_w),
|
||||
lambda b, c, w: te.sum(
|
||||
data_pad[b, dc, w + dw].astype(out_dtype)
|
||||
* transformed_kernel[c, dc, dw].astype(out_dtype),
|
||||
axis=[dc, dw],
|
||||
),
|
||||
tag="conv1d_transpose_ncw",
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def group_conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding, groups):
|
||||
"""Transposed 1D group convolution ncw forward operator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
data : tvm.te.Tensor
|
||||
3-D with shape [batch, in_channel, in_width]
|
||||
|
||||
kernel : tvm.te.Tensor
|
||||
3-D with shape [in_channel, num_filter, filter_width]
|
||||
|
||||
stride : ints
|
||||
The spatial stride along width
|
||||
|
||||
padding : int or str
|
||||
Padding size, or ['VALID', 'SAME']
|
||||
|
||||
out_dtype : str
|
||||
The output data type. This is used for mixed precision.
|
||||
|
||||
output_padding : ints
|
||||
Used to recover the actual output shape in case there are more
|
||||
than one possible shape. Must be smaller than stride.
|
||||
|
||||
groups : int
|
||||
number of groups
|
||||
|
||||
Returns
|
||||
-------
|
||||
output : tvm.te.Tensor
|
||||
3-D with shape [batch, out_channel, out_width]
|
||||
|
||||
"""
|
||||
if groups == 1:
|
||||
return conv1d_transpose_ncw(data, kernel, stride, padding, out_dtype, output_padding)
|
||||
|
||||
_, in_channels, _ = data.shape
|
||||
|
||||
assert in_channels % groups == 0, (
|
||||
f"input channels {in_channels} must divide group size {groups}"
|
||||
)
|
||||
|
||||
data_pad, transformed_kernel = _conv1d_transpose_ncw_preprocess(
|
||||
data, kernel, stride, padding, out_dtype, output_padding
|
||||
)
|
||||
|
||||
batch, in_channels, in_w = data_pad.shape
|
||||
out_c, _, filter_w = transformed_kernel.shape
|
||||
|
||||
# convolution stage
|
||||
out_channels = simplify(out_c * groups)
|
||||
out_w = simplify(in_w - filter_w + 1)
|
||||
dc = te.reduce_axis((0, in_channels // groups), name="dc")
|
||||
dw = te.reduce_axis((0, filter_w), name="dw")
|
||||
|
||||
# data: batch, in_channels, out_w
|
||||
# weight: out_channels // G, in_channels, out_w
|
||||
return te.compute(
|
||||
(batch, out_channels, out_w),
|
||||
lambda b, c, w: te.sum(
|
||||
data_pad[
|
||||
b, c // (out_channels // groups) * (in_channels // groups) + dc, w + dw
|
||||
].astype(out_dtype)
|
||||
* transformed_kernel[
|
||||
c % (out_channels // groups),
|
||||
c // (out_channels // groups) * (in_channels // groups) + dc,
|
||||
dw,
|
||||
].astype(out_dtype),
|
||||
axis=[dc, dw],
|
||||
),
|
||||
tag="group_conv1d_transpose_ncw",
|
||||
)
|
||||
Reference in New Issue
Block a user