Construct an ConvolutionModule object. Args: channels (int): The number of channels of conv layers. kernel_size (int): Kernel size of conv layers. causal (int): Whether use causal convolution or not
(
self,
channels: int,
kernel_size: int = 15,
activation: nn.Module = nn.ReLU(),
norm: str = "batch_norm",
causal: bool = False,
bias: bool = True,
)
| 8 | """ConvolutionModule in Conformer model.""" |
| 9 | |
| 10 | def __init__( |
| 11 | self, |
| 12 | channels: int, |
| 13 | kernel_size: int = 15, |
| 14 | activation: nn.Module = nn.ReLU(), |
| 15 | norm: str = "batch_norm", |
| 16 | causal: bool = False, |
| 17 | bias: bool = True, |
| 18 | ): |
| 19 | """Construct an ConvolutionModule object. |
| 20 | Args: |
| 21 | channels (int): The number of channels of conv layers. |
| 22 | kernel_size (int): Kernel size of conv layers. |
| 23 | causal (int): Whether use causal convolution or not |
| 24 | """ |
| 25 | super().__init__() |
| 26 | |
| 27 | self.pointwise_conv1 = nn.Conv1d( |
| 28 | channels, |
| 29 | 2 * channels, |
| 30 | kernel_size=1, |
| 31 | stride=1, |
| 32 | padding=0, |
| 33 | bias=bias, |
| 34 | ) |
| 35 | # self.lorder is used to distinguish if it's a causal convolution, |
| 36 | # if self.lorder > 0: it's a causal convolution, the input will be |
| 37 | # padded with self.lorder frames on the left in forward. |
| 38 | # else: it's a symmetrical convolution |
| 39 | if causal: |
| 40 | padding = 0 |
| 41 | self.lorder = kernel_size - 1 |
| 42 | else: |
| 43 | # kernel_size should be an odd number for none causal convolution |
| 44 | assert (kernel_size - 1) % 2 == 0 |
| 45 | padding = (kernel_size - 1) // 2 |
| 46 | self.lorder = 0 |
| 47 | self.depthwise_conv = nn.Conv1d( |
| 48 | channels, |
| 49 | channels, |
| 50 | kernel_size, |
| 51 | stride=1, |
| 52 | padding=padding, |
| 53 | groups=channels, |
| 54 | bias=bias, |
| 55 | ) |
| 56 | |
| 57 | assert norm in ["batch_norm", "layer_norm"] |
| 58 | if norm == "batch_norm": |
| 59 | self.use_layer_norm = False |
| 60 | self.norm = nn.BatchNorm1d(channels) |
| 61 | else: |
| 62 | self.use_layer_norm = True |
| 63 | self.norm = nn.LayerNorm(channels) |
| 64 | |
| 65 | self.pointwise_conv2 = nn.Conv1d( |
| 66 | channels, |
| 67 | channels, |