|
import numpy as np |
|
from torch import nn |
|
from torch.nn.utils import weight_norm |
|
|
|
|
|
class MelganDiscriminator(nn.Module): |
|
def __init__( |
|
self, |
|
in_channels=1, |
|
out_channels=1, |
|
kernel_sizes=(5, 3), |
|
base_channels=16, |
|
max_channels=1024, |
|
downsample_factors=(4, 4, 4, 4), |
|
groups_denominator=4, |
|
): |
|
super().__init__() |
|
self.layers = nn.ModuleList() |
|
|
|
layer_kernel_size = np.prod(kernel_sizes) |
|
layer_padding = (layer_kernel_size - 1) // 2 |
|
|
|
|
|
self.layers += [ |
|
nn.Sequential( |
|
nn.ReflectionPad1d(layer_padding), |
|
weight_norm(nn.Conv1d(in_channels, base_channels, layer_kernel_size, stride=1)), |
|
nn.LeakyReLU(0.2, inplace=True), |
|
) |
|
] |
|
|
|
|
|
layer_in_channels = base_channels |
|
for downsample_factor in downsample_factors: |
|
layer_out_channels = min(layer_in_channels * downsample_factor, max_channels) |
|
layer_kernel_size = downsample_factor * 10 + 1 |
|
layer_padding = (layer_kernel_size - 1) // 2 |
|
layer_groups = layer_in_channels // groups_denominator |
|
self.layers += [ |
|
nn.Sequential( |
|
weight_norm( |
|
nn.Conv1d( |
|
layer_in_channels, |
|
layer_out_channels, |
|
kernel_size=layer_kernel_size, |
|
stride=downsample_factor, |
|
padding=layer_padding, |
|
groups=layer_groups, |
|
) |
|
), |
|
nn.LeakyReLU(0.2, inplace=True), |
|
) |
|
] |
|
layer_in_channels = layer_out_channels |
|
|
|
|
|
layer_padding1 = (kernel_sizes[0] - 1) // 2 |
|
layer_padding2 = (kernel_sizes[1] - 1) // 2 |
|
self.layers += [ |
|
nn.Sequential( |
|
weight_norm( |
|
nn.Conv1d( |
|
layer_out_channels, |
|
layer_out_channels, |
|
kernel_size=kernel_sizes[0], |
|
stride=1, |
|
padding=layer_padding1, |
|
) |
|
), |
|
nn.LeakyReLU(0.2, inplace=True), |
|
), |
|
weight_norm( |
|
nn.Conv1d( |
|
layer_out_channels, out_channels, kernel_size=kernel_sizes[1], stride=1, padding=layer_padding2 |
|
) |
|
), |
|
] |
|
|
|
def forward(self, x): |
|
feats = [] |
|
for layer in self.layers: |
|
x = layer(x) |
|
feats.append(x) |
|
return x, feats |
|
|
