ConvolutionModule in Conformer model.
| 23 | |
| 24 | |
| 25 | class ConvolutionModule(nn.Module): |
| 26 | """ConvolutionModule in Conformer model.""" |
| 27 | |
| 28 | def __init__( |
| 29 | self, |
| 30 | channels: int, |
| 31 | kernel_size: int = 15, |
| 32 | activation: nn.Module = nn.ReLU(), |
| 33 | norm: str = "batch_norm", |
| 34 | causal: bool = False, |
| 35 | bias: bool = True, |
| 36 | norm_eps: float = 1e-5, |
| 37 | ): |
| 38 | """Construct an ConvolutionModule object. |
| 39 | Args: |
| 40 | channels (int): The number of channels of conv layers. |
| 41 | kernel_size (int): Kernel size of conv layers. |
| 42 | causal (int): Whether use causal convolution or not |
| 43 | """ |
| 44 | super().__init__() |
| 45 | |
| 46 | self.pointwise_conv1 = nn.Conv1d( |
| 47 | channels, |
| 48 | 2 * channels, |
| 49 | kernel_size=1, |
| 50 | stride=1, |
| 51 | padding=0, |
| 52 | bias=bias, |
| 53 | ) |
| 54 | # self.lorder is used to distinguish if it's a causal convolution, |
| 55 | # if self.lorder > 0: it's a causal convolution, the input will be |
| 56 | # padded with self.lorder frames on the left in forward. |
| 57 | # else: it's a symmetrical convolution |
| 58 | if causal: |
| 59 | padding = 0 |
| 60 | self.lorder = kernel_size - 1 |
| 61 | else: |
| 62 | # kernel_size should be an odd number for none causal convolution |
| 63 | assert (kernel_size - 1) % 2 == 0 |
| 64 | padding = (kernel_size - 1) // 2 |
| 65 | self.lorder = 0 |
| 66 | self.depthwise_conv = nn.Conv1d( |
| 67 | channels, |
| 68 | channels, |
| 69 | kernel_size, |
| 70 | stride=1, |
| 71 | padding=padding, |
| 72 | groups=channels, |
| 73 | bias=bias, |
| 74 | ) |
| 75 | |
| 76 | assert norm in ['batch_norm', 'layer_norm', 'rms_norm'] |
| 77 | if norm == "batch_norm": |
| 78 | self.use_layer_norm = False |
| 79 | self.norm = WENET_NORM_CLASSES['batch_norm'](channels, |
| 80 | eps=norm_eps) |
| 81 | else: |
| 82 | self.use_layer_norm = True |