# 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", )