def _get_conv(
self,
in_channels,
out_channels,
kernel_size=11,
stride=1,
dilation=1,
padding=0,
bias=False,
groups=1,
heads=-1,
separable=False,
quantize=False,
):
use_mask = self.conv_mask
if use_mask:
return MaskedConv1d(
in_channels,
out_channels,
kernel_size,
stride=stride,
dilation=dilation,
padding=padding,
bias=bias,
groups=groups,
heads=heads,
use_mask=use_mask,
quantize=quantize,
)
else:
if PYTORCH_QUANTIZATION_AVAILABLE and quantize:
return quant_nn.QuantConv1d(
in_channels,
out_channels,
kernel_size,
stride=stride,
dilation=dilation,
padding=padding,
bias=bias,
groups=groups,
)
elif not PYTORCH_QUANTIZATION_AVAILABLE and quantize:
raise ImportError(
"pytorch-quantization is not installed. Install from "
"https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
)
else:
return nn.Conv1d(
in_channels,
out_channels,
kernel_size,
stride=stride,
dilation=dilation,
padding=padding,
bias=bias,
groups=groups,
)
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
groups=1,
heads=-1,
bias=False,
use_mask=True,
quantize=False,
):
super(MaskedConv1d, self).__init__()
if not (heads == -1 or groups == in_channels):
raise ValueError("Only use heads for depthwise convolutions")
self.real_out_channels = out_channels
if heads != -1:
in_channels = heads
out_channels = heads
groups = heads
if PYTORCH_QUANTIZATION_AVAILABLE and quantize:
self.conv = quant_nn.QuantConv1d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias,
)
elif not PYTORCH_QUANTIZATION_AVAILABLE and quantize:
raise ImportError(
"pytorch-quantization is not installed. Install from "
"https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
)
else:
self.conv = nn.Conv1d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias,
)
self.use_mask = use_mask
self.heads = heads
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
groups=1,
heads=-1,
bias=False,
use_mask=True,
quantize=False,
):
super(MaskedConv1d, self).__init__()
if not (heads == -1 or groups == in_channels):
raise ValueError("Only use heads for depthwise convolutions")
self.real_out_channels = out_channels
if heads != -1:
in_channels = heads
out_channels = heads
groups = heads
# preserve original padding
self._padding = padding
# if padding is a tuple/list, it is considered as asymmetric padding
if type(padding) in (tuple, list):
self.pad_layer = nn.ConstantPad1d(padding, value=0.0)
# reset padding for conv since pad_layer will handle this
padding = 0
else:
self.pad_layer = None
if PYTORCH_QUANTIZATION_AVAILABLE and quantize:
self.conv = quant_nn.QuantConv1d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias,
)
elif not PYTORCH_QUANTIZATION_AVAILABLE and quantize:
raise ImportError(
"pytorch-quantization is not installed. Install from "
"https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization."
)
else:
self.conv = nn.Conv1d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias,
)
self.use_mask = use_mask
self.heads = heads
# Calculations for "same" padding cache
self.same_padding = (self.conv.stride[0] == 1) and (
2 * self.conv.padding[0] == self.conv.dilation[0] * (self.conv.kernel_size[0] - 1)
)
if self.pad_layer is None:
self.same_padding_asymmetric = False
else:
self.same_padding_asymmetric = (self.conv.stride[0] == 1) and (
sum(self._padding) == self.conv.dilation[0] * (self.conv.kernel_size[0] - 1)
)
# `self.lens` caches consecutive integers from 0 to `self.max_len` that are used to compute the mask for a
# batch. Recomputed to bigger size as needed. Stored on a device of the latest batch lens.
if self.use_mask:
self.max_len = 0
self.lens = None
