|
|
|
|
|
""" |
|
References: |
|
- https://github.com/jik876/hifi-gan |
|
- https://github.com/kan-bayashi/ParallelWaveGAN |
|
""" |
|
|
|
import torch |
|
|
|
|
|
class Conv1d(torch.nn.Conv1d): |
|
""" |
|
Conv1d module with customized initialization. |
|
""" |
|
|
|
def __init__(self, *args, **kwargs): |
|
super(Conv1d, self).__init__(*args, **kwargs) |
|
|
|
def reset_parameters(self): |
|
torch.nn.init.kaiming_normal_(self.weight, nonlinearity="relu") |
|
if self.bias is not None: |
|
torch.nn.init.constant_(self.bias, 0.0) |
|
|
|
|
|
class Conv1d1x1(Conv1d): |
|
""" |
|
1x1 Conv1d with customized initialization. |
|
""" |
|
|
|
def __init__(self, in_channels, out_channels, bias): |
|
super(Conv1d1x1, self).__init__(in_channels, out_channels, kernel_size=1, padding=0, dilation=1, bias=bias) |
|
|
|
|
|
class HiFiGANResidualBlock(torch.nn.Module): |
|
"""Residual block module in HiFiGAN.""" |
|
|
|
def __init__(self, |
|
kernel_size=3, |
|
channels=512, |
|
dilations=(1, 3, 5), |
|
bias=True, |
|
use_additional_convs=True, |
|
nonlinear_activation="LeakyReLU", |
|
nonlinear_activation_params={"negative_slope": 0.1}, ): |
|
""" |
|
Initialize HiFiGANResidualBlock module. |
|
|
|
Args: |
|
kernel_size (int): Kernel size of dilation convolution layer. |
|
channels (int): Number of channels for convolution layer. |
|
dilations (List[int]): List of dilation factors. |
|
use_additional_convs (bool): Whether to use additional convolution layers. |
|
bias (bool): Whether to add bias parameter in convolution layers. |
|
nonlinear_activation (str): Activation function module name. |
|
nonlinear_activation_params (dict): Hyperparameters for activation function. |
|
""" |
|
super().__init__() |
|
self.use_additional_convs = use_additional_convs |
|
self.convs1 = torch.nn.ModuleList() |
|
if use_additional_convs: |
|
self.convs2 = torch.nn.ModuleList() |
|
assert kernel_size % 2 == 1, "Kernel size must be odd number." |
|
for dilation in dilations: |
|
self.convs1 += [torch.nn.Sequential(getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params), |
|
torch.nn.Conv1d(channels, |
|
channels, |
|
kernel_size, |
|
1, |
|
dilation=dilation, |
|
bias=bias, |
|
padding=(kernel_size - 1) // 2 * dilation, ), )] |
|
if use_additional_convs: |
|
self.convs2 += [torch.nn.Sequential(getattr(torch.nn, nonlinear_activation)(**nonlinear_activation_params), |
|
torch.nn.Conv1d(channels, |
|
channels, |
|
kernel_size, |
|
1, |
|
dilation=1, |
|
bias=bias, |
|
padding=(kernel_size - 1) // 2, ), )] |
|
|
|
def forward(self, x): |
|
""" |
|
Calculate forward propagation. |
|
|
|
Args: |
|
x (Tensor): Input tensor (B, channels, T). |
|
|
|
Returns: |
|
Tensor: Output tensor (B, channels, T). |
|
""" |
|
for idx in range(len(self.convs1)): |
|
xt = self.convs1[idx](x) |
|
if self.use_additional_convs: |
|
xt = self.convs2[idx](xt) |
|
x = xt + x |
|
return x |
|
|