def __init__(self,
encoder_dim: int = 512,
expansion_factor: int = 4,
dropout_p: float = 0.1,
device: torch.device = 'cuda') -> None:
super(FeedForwardModule, self).__init__()
self.device = device
self.sequential = nn.Sequential(
LayerNorm(encoder_dim),
Linear(encoder_dim, encoder_dim * expansion_factor, bias=True),
Swish(),
nn.Dropout(p=dropout_p),
Linear(encoder_dim * expansion_factor, encoder_dim, bias=True),
nn.Dropout(p=dropout_p),
)
def __init__(
self,
in_channels: int,
kernel_size: int = 31,
expansion_factor: int = 2,
dropout_p: float = 0.1,
device: torch.device = 'cuda',
) -> None:
super(ConformerConvModule, self).__init__()
assert (
kernel_size - 1
) % 2 == 0, "kernel_size should be a odd number for 'SAME' padding"
assert expansion_factor == 2, "Currently, Only Supports expansion_factor 2"
self.device = device
self.sequential = nn.Sequential(
LayerNorm(in_channels),
Transpose(shape=(1, 2)),
PointwiseConv1d(in_channels,
in_channels * expansion_factor,
stride=1,
padding=0,
bias=True),
GLU(dim=1),
DepthwiseConv1d(in_channels,
in_channels,
kernel_size,
stride=1,
padding=(kernel_size - 1) // 2),
nn.BatchNorm1d(in_channels),
Swish(),
PointwiseConv1d(in_channels,
in_channels,
stride=1,
padding=0,
bias=True),
nn.Dropout(p=dropout_p),
)
class Conv2dExtractor(nn.Module):
"""
Provides inteface of convolutional extractor.
Note:
Do not use this class directly, use one of the sub classes.
Define the 'self.conv' class variable.
Inputs: inputs, input_lengths
- **inputs** (batch, time, dim): Tensor containing input vectors
- **input_lengths**: Tensor containing containing sequence lengths
Returns: outputs, output_lengths
- **outputs**: Tensor produced by the convolution
- **output_lengths**: Tensor containing sequence lengths produced by the convolution
"""
supported_activations = {
'hardtanh': nn.Hardtanh(0, 20, inplace=True),
'relu': nn.ReLU(inplace=True),
'elu': nn.ELU(inplace=True),
'leaky_relu': nn.LeakyReLU(inplace=True),
'gelu': nn.GELU(),
'swish': Swish(),
}
def __init__(self, input_dim: int, activation: str = 'hardtanh') -> None:
super(Conv2dExtractor, self).__init__()
self.input_dim = input_dim
self.activation = Conv2dExtractor.supported_activations[activation]
self.conv = None
def get_output_lengths(self, seq_lengths: Tensor):
assert self.conv is not None, "self.conv should be defined"
for module in self.conv:
if isinstance(module, nn.Conv2d):
numerator = seq_lengths + 2 * module.padding[
1] - module.dilation[1] * (module.kernel_size[1] - 1) - 1
seq_lengths = numerator.float() / float(module.stride[1])
seq_lengths = seq_lengths.int() + 1
elif isinstance(module, nn.MaxPool2d):
seq_lengths >>= 1
return seq_lengths.int()
def get_output_dim(self):
if isinstance(self, VGGExtractor):
output_dim = (self.input_dim - 1
) << 5 if self.input_dim % 2 else self.input_dim << 5
elif isinstance(self, DeepSpeech2Extractor):
output_dim = int(math.floor(self.input_dim + 2 * 20 - 41) / 2 + 1)
output_dim = int(math.floor(output_dim + 2 * 10 - 21) / 2 + 1)
output_dim <<= 5
elif isinstance(self, Conv2dSubsampling):
factor = ((self.input_dim - 1) // 2 - 1) // 2
output_dim = self.out_channels * factor
else:
raise ValueError(f"Unsupported Extractor : {self.extractor}")
return output_dim
def forward(self, inputs: Tensor,
input_lengths: Tensor) -> Tuple[Tensor, Tensor]:
"""
inputs: torch.FloatTensor (batch, time, dimension)
input_lengths: torch.IntTensor (batch)
"""
outputs, output_lengths = self.conv(
inputs.unsqueeze(1).transpose(2, 3), input_lengths)
batch_size, channels, dimension, seq_lengths = outputs.size()
outputs = outputs.permute(0, 3, 1, 2)
outputs = outputs.view(batch_size, seq_lengths, channels * dimension)
return outputs, output_lengths