Source code for espnet2.asr_transducer.encoder.blocks.conv_input

"""ConvInput block for Transducer encoder."""

from typing import Optional, Tuple, Union

import torch

from espnet2.asr_transducer.utils import get_convinput_module_parameters

[docs]class ConvInput(torch.nn.Module): """ConvInput module definition. Args: input_size: Input size. conv_size: Convolution size. subsampling_factor: Subsampling factor. vgg_like: Whether to use a VGG-like network. output_size: Block output dimension. """ def __init__( self, input_size: int, conv_size: Union[int, Tuple], subsampling_factor: int = 4, vgg_like: bool = True, output_size: Optional[int] = None, ) -> None: """Construct a ConvInput object.""" super().__init__() self.subsampling_factor = subsampling_factor self.vgg_like = vgg_like if vgg_like: conv_size1, conv_size2 = conv_size self.maxpool_kernel1, output_proj = get_convinput_module_parameters( input_size, conv_size2, subsampling_factor, is_vgg=True ) self.conv = torch.nn.Sequential( torch.nn.Conv2d(1, conv_size1, 3, stride=1, padding=1), torch.nn.ReLU(), torch.nn.Conv2d(conv_size1, conv_size1, 3, stride=1, padding=0), torch.nn.ReLU(), torch.nn.MaxPool2d( self.maxpool_kernel1, stride=2, padding=0, ceil_mode=True ), torch.nn.Conv2d(conv_size1, conv_size2, 3, stride=1, padding=1), torch.nn.ReLU(), torch.nn.Conv2d(conv_size2, conv_size2, 3, stride=1, padding=0), torch.nn.ReLU(), torch.nn.MaxPool2d(2, stride=2, padding=0, ceil_mode=True), ) else: ( self.conv_kernel2, self.conv_stride2, ), output_proj = get_convinput_module_parameters( input_size, conv_size, subsampling_factor, is_vgg=False ) self.conv = torch.nn.Sequential( torch.nn.Conv2d(1, conv_size, 3, 2), torch.nn.ReLU(), torch.nn.Conv2d( conv_size, conv_size, self.conv_kernel2, self.conv_stride2 ), torch.nn.ReLU(), ) self.min_frame_length = 7 if subsampling_factor < 6 else 11 if output_size is not None: self.output = torch.nn.Linear(output_proj, output_size) self.output_size = output_size else: self.output = None self.output_size = output_proj
[docs] def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> Tuple[torch.Tensor, torch.Tensor]: """Encode input sequences. Args: x: ConvInput input sequences. (B, T, D_feats) mask: Mask of input sequences. (B, 1, T) Returns: x: ConvInput output sequences. (B, sub(T), D_out) mask: Mask of output sequences. (B, 1, sub(T)) """ x = self.conv(x.unsqueeze(1)) b, c, t, f = x.size() x = x.transpose(1, 2).contiguous().view(b, t, c * f) if self.output is not None: x = self.output(x) if mask is not None: mask = mask[:, : x.size(1)] return x, mask