wehub-resource-sync
247 lines
8.4 KiB
Python
247 lines
8.4 KiB
Python
# 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
|
|
# ruff: noqa: F821
|
|
"""Transposed 2D convolution operators (sometimes called Deconvolution)."""
|
|
|
|
import collections
|
|
|
|
from tvm import te
|
|
|
|
from ..utils import simplify
|
|
from .dilate import dilate
|
|
from .pad import pad
|
|
from .utils import get_pad_tuple
|
|
|
|
|
|
def _ntuple(n):
|
|
def parse(x):
|
|
if isinstance(x, collections.abc.Iterable):
|
|
assert len(x) == n, f"Input can only have {n} elements, but got {len(x)} instead: {x}."
|
|
return x
|
|
return tuple(repeat(x, n))
|
|
|
|
return parse
|
|
|
|
|
|
_single = _ntuple(1)
|
|
_pair = _ntuple(2)
|
|
_triple = _ntuple(3)
|
|
_quadruple = _ntuple(4)
|
|
|
|
|
|
def conv2d_transpose_nchw(Input, Filter, strides, padding, out_dtype, output_padding):
|
|
"""Transposed 2D convolution nchw forward operator.
|
|
|
|
Parameters
|
|
----------
|
|
Input : tvm.te.Tensor
|
|
4-D with shape [batch, in_channel, in_height, in_width]
|
|
|
|
Filter : tvm.te.Tensor
|
|
4-D with shape [in_channel, num_filter, filter_height, filter_width]
|
|
|
|
strides : tuple of two ints
|
|
The spatial stride along height and 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 : tuple of ints
|
|
Used to get the right output shape for gradients
|
|
|
|
Returns
|
|
-------
|
|
Output : tvm.te.Tensor
|
|
4-D with shape [batch, out_channel, out_height, out_width]
|
|
"""
|
|
return declaration_conv2d_transpose_impl(
|
|
Input, Filter, strides, padding, out_dtype, output_padding=output_padding
|
|
)
|
|
|
|
|
|
def conv2d_transpose_nchw_preprocess(data, kernel, strides, padding, out_dtype, output_padding):
|
|
"""Preprocess data and kernel to make the compute pattern
|
|
of conv2d_transpose the same as conv2d"""
|
|
batch, in_c, in_h, in_w = data.shape
|
|
_, out_c, filter_h, filter_w = kernel.shape
|
|
stride_h, stride_w = strides
|
|
opad_h, opad_w = output_padding
|
|
assert opad_h < stride_h and opad_w < stride_w
|
|
# dilate data
|
|
data_dilate = dilate(data, [1, 1, stride_h, stride_w], name="data_dilate")
|
|
# pad data
|
|
fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
|
|
bpad_top = filter_h - 1 - fpad_top
|
|
bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
|
|
bpad_left = filter_w - 1 - fpad_left
|
|
bpad_right = filter_w - 1 - fpad_right + opad_w
|
|
data_pad = pad(
|
|
data_dilate, [0, 0, bpad_top, bpad_left], [0, 0, bpad_bottom, bpad_right], name="data_pad"
|
|
)
|
|
# transform kernel layout from IOHW to OIHW, and rotate kernel by 180 degrees
|
|
kernel_transform = te.compute(
|
|
(out_c, in_c, filter_h, filter_w),
|
|
lambda o, i, h, w: kernel[i][o][filter_h - 1 - h][filter_w - 1 - w],
|
|
name="kernel_transform",
|
|
)
|
|
return data_pad, kernel_transform
|
|
|
|
|
|
def declaration_conv2d_transpose_impl(data, kernel, strides, padding, out_dtype, output_padding):
|
|
"""Implementation of conv2d transpose"""
|
|
data_pad, kernel_transform = conv2d_transpose_nchw_preprocess(
|
|
data, kernel, strides, padding, out_dtype, output_padding
|
|
)
|
|
batch, in_c, in_h, in_w = data_pad.shape
|
|
out_c, _, filter_h, filter_w = kernel_transform.shape
|
|
|
|
# convolution stage
|
|
out_c = simplify(out_c)
|
|
|
|
out_h = simplify(in_h - filter_h + 1)
|
|
out_w = simplify(in_w - filter_w + 1)
|
|
dc = te.reduce_axis((0, in_c), name="dc")
|
|
dh = te.reduce_axis((0, filter_h), name="dh")
|
|
dw = te.reduce_axis((0, filter_w), name="dw")
|
|
|
|
Output = te.compute(
|
|
(batch, out_c, out_h, out_w),
|
|
lambda b, c, h, w: te.sum(
|
|
data_pad[b, dc, h + dh, w + dw].astype(out_dtype)
|
|
* kernel_transform[c, dc, dh, dw].astype(out_dtype),
|
|
axis=[dc, dh, dw],
|
|
),
|
|
tag="conv2d_transpose_nchw",
|
|
)
|
|
|
|
return Output
|
|
|
|
|
|
def group_conv2d_transpose_nchw(data, kernel, stride, padding, out_dtype, output_padding, groups):
|
|
"""Group convolution operator in NCHW layout.
|
|
|
|
Parameters
|
|
----------
|
|
data : tvm.te.Tensor
|
|
4-D with shape [batch, in_channel, in_height, in_width]
|
|
|
|
kernel : tvm.te.Tensor
|
|
4-D with shape [in_channel, out_channel // groups, filter_height, filter_width]
|
|
|
|
stride : int or a list/tuple of two ints
|
|
Stride size, or [stride_height, stride_width]
|
|
|
|
padding : int or a list/tuple of 2 or 4 ints
|
|
padding size, or
|
|
[pad_height, pad_width] for 2 ints, or
|
|
[pad_top, pad_left, pad_bottom, pad_right] for 4 ints
|
|
|
|
out_dtype : str
|
|
The output data type. This is used for mixed precision.
|
|
|
|
output_padding : tuple of ints
|
|
Used to get the right output shape for gradients
|
|
|
|
groups : int
|
|
number of groups
|
|
|
|
out_dtype : str
|
|
The output type. This is used for mixed precision.
|
|
|
|
Returns
|
|
-------
|
|
Output : tvm.te.Tensor
|
|
4-D with shape [batch, out_channel, out_height, out_width]
|
|
"""
|
|
if groups == 1:
|
|
return conv2d_transpose_nchw(data, kernel, stride, padding, out_dtype, output_padding)
|
|
|
|
# some pre-processing and prelimnary checks
|
|
if out_dtype is None:
|
|
out_dtype = data.dtype
|
|
|
|
batch, in_channels, in_h, in_w = data.shape
|
|
_, out_c, filter_h, filter_w = kernel.shape
|
|
assert in_channels % groups == 0, (
|
|
f"input channels {in_channels} must divide group size {groups}"
|
|
)
|
|
# assert out_c % groups == 0, f"output channels {in_c} must divide group size {groups}"
|
|
|
|
strides = _pair(stride)
|
|
# padding = _pair(padding)
|
|
# output_padding = _pair(output_padding)
|
|
# dilation = _pair(dilation)
|
|
|
|
stride_h, stride_w = strides
|
|
opad_h, opad_w = output_padding
|
|
assert opad_h < stride_h and opad_w < stride_w, (
|
|
f"[{output_padding}] opad_h:{opad_h} < stride_h:{stride_h} \
|
|
and opad_w:{opad_w} < stride_w:{stride_w} does not satisfy."
|
|
)
|
|
# dilate data
|
|
data_dilate = dilate(data, [1, 1, stride_h, stride_w], name="data_dilate")
|
|
# pad data
|
|
fpad_top, fpad_left, fpad_bottom, fpad_right = get_pad_tuple(padding, (filter_h, filter_w))
|
|
bpad_top = filter_h - 1 - fpad_top
|
|
bpad_bottom = filter_h - 1 - fpad_bottom + opad_h
|
|
bpad_left = filter_w - 1 - fpad_left
|
|
bpad_right = filter_w - 1 - fpad_right + opad_w
|
|
data_pad = pad(
|
|
data_dilate, [0, 0, bpad_top, bpad_left], [0, 0, bpad_bottom, bpad_right], name="data_pad"
|
|
)
|
|
# transform kernel layout from IOHW to OIHW, and rotate kernel by 180 degrees
|
|
kernel_transform = te.compute(
|
|
(out_c, in_channels, filter_h, filter_w),
|
|
lambda i, o, h, w: kernel[o][i][filter_h - 1 - h][filter_w - 1 - w],
|
|
name="kernel_transform",
|
|
)
|
|
|
|
batch, in_channels, in_h, in_w = data_pad.shape
|
|
out_c, _, filter_h, filter_w = kernel_transform.shape
|
|
|
|
# convolution stage
|
|
out_channels = simplify(out_c * groups)
|
|
|
|
out_h = simplify(in_h - filter_h + 1)
|
|
out_w = simplify(in_w - filter_w + 1)
|
|
dc = te.reduce_axis((0, in_channels // groups), name="dc")
|
|
dh = te.reduce_axis((0, filter_h), name="dh")
|
|
dw = te.reduce_axis((0, filter_w), name="dw")
|
|
|
|
# data: batch, in_channels, out_h, out_w
|
|
# weight: out_channels // G, in_channels, out_h, out_w
|
|
return te.compute(
|
|
(batch, out_channels, out_h, out_w),
|
|
lambda b, c, h, w: te.sum(
|
|
data_pad[
|
|
b, c // (out_channels // groups) * (in_channels // groups) + dc, h + dh, w + dw
|
|
].astype(out_dtype)
|
|
* kernel_transform[
|
|
c % (out_channels // groups),
|
|
c // (out_channels // groups) * (in_channels // groups) + dc,
|
|
dh,
|
|
dw,
|
|
].astype(out_dtype),
|
|
axis=[dc, dh, dw],
|
|
),
|
|
tag="group_conv2d_transpose_nchw",
|
|
)
|