| 148 | |
| 149 | class ConvReluNorm(nn.Module): |
| 150 | def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout): |
| 151 | super().__init__() |
| 152 | self.in_channels = in_channels |
| 153 | self.hidden_channels = hidden_channels |
| 154 | self.out_channels = out_channels |
| 155 | self.kernel_size = kernel_size |
| 156 | self.n_layers = n_layers |
| 157 | self.p_dropout = p_dropout |
| 158 | assert n_layers > 1, "Number of layers should be larger than 0." |
| 159 | |
| 160 | self.conv_layers = nn.ModuleList() |
| 161 | self.norm_layers = nn.ModuleList() |
| 162 | self.conv_layers.append(nn.Conv1d(in_channels, hidden_channels, kernel_size, padding=kernel_size // 2)) |
| 163 | self.norm_layers.append(LayerNorm(hidden_channels)) |
| 164 | self.relu_drop = nn.Sequential( |
| 165 | nn.ReLU(), |
| 166 | nn.Dropout(p_dropout)) |
| 167 | for _ in range(n_layers - 1): |
| 168 | self.conv_layers.append(nn.Conv1d(hidden_channels, hidden_channels, kernel_size, padding=kernel_size // 2)) |
| 169 | self.norm_layers.append(LayerNorm(hidden_channels)) |
| 170 | self.proj = nn.Conv1d(hidden_channels, out_channels, 1) |
| 171 | self.proj.weight.data.zero_() |
| 172 | self.proj.bias.data.zero_() |
| 173 | |
| 174 | def forward(self, x, x_mask): |
| 175 | x_org = x |