MelGAN discriminator module.
| 192 | |
| 193 | |
| 194 | class MelGANDiscriminator(torch.nn.Module): |
| 195 | """MelGAN discriminator module.""" |
| 196 | |
| 197 | def __init__(self, |
| 198 | in_channels=1, |
| 199 | out_channels=1, |
| 200 | kernel_sizes=[5, 3], |
| 201 | channels=16, |
| 202 | max_downsample_channels=1024, |
| 203 | bias=True, |
| 204 | downsample_scales=[4, 4, 4, 4], |
| 205 | nonlinear_activation="LeakyReLU", |
| 206 | nonlinear_activation_params={"negative_slope": 0.2}, |
| 207 | pad="ReflectionPad1d", |
| 208 | pad_params={}, |
| 209 | ): |
| 210 | """Initilize MelGAN discriminator module. |
| 211 | |
| 212 | Args: |
| 213 | in_channels (int): Number of input channels. |
| 214 | out_channels (int): Number of output channels. |
| 215 | kernel_sizes (list): List of two kernel sizes. The prod will be used for the first conv layer, |
| 216 | and the first and the second kernel sizes will be used for the last two layers. |
| 217 | For example if kernel_sizes = [5, 3], the first layer kernel size will be 5 * 3 = 15, |
| 218 | the last two layers' kernel size will be 5 and 3, respectively. |
| 219 | channels (int): Initial number of channels for conv layer. |
| 220 | max_downsample_channels (int): Maximum number of channels for downsampling layers. |
| 221 | bias (bool): Whether to add bias parameter in convolution layers. |
| 222 | downsample_scales (list): List of downsampling scales. |
| 223 | nonlinear_activation (str): Activation function module name. |
| 224 | nonlinear_activation_params (dict): Hyperparameters for activation function. |
| 225 | pad (str): Padding function module name before dilated convolution layer. |
| 226 | pad_params (dict): Hyperparameters for padding function. |
| 227 | |
| 228 | """ |
| 229 | super(MelGANDiscriminator, self).__init__() |
| 230 | self.layers = torch.nn.ModuleList() |
| 231 | |
| 232 | # check kernel size is valid |
| 233 | assert len(kernel_sizes) == 2 |
| 234 | assert kernel_sizes[0] % 2 == 1 |
| 235 | assert kernel_sizes[1] % 2 == 1 |
| 236 | |
| 237 | # add first layer |
| 238 | self.layers += [ |
| 239 | torch.nn.Sequential( |
| 240 | getattr(torch.nn, pad)((np.prod(kernel_sizes) - 1) // 2, **pad_params), |
| 241 | torch.nn.Conv1d(in_channels, channels, np.prod(kernel_sizes), bias=bias), |
| 242 | getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params), |
| 243 | ) |
| 244 | ] |
| 245 | |
| 246 | # add downsample layers |
| 247 | in_chs = channels |
| 248 | for downsample_scale in downsample_scales: |
| 249 | out_chs = min(in_chs * downsample_scale, max_downsample_channels) |
| 250 | self.layers += [ |
| 251 | torch.nn.Sequential( |