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