138 lines
4.0 KiB
Python
138 lines
4.0 KiB
Python
# Copyright (c) 2023 Amphion.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.nn import functional as F
|
|
from .SubLayers import MultiHeadAttention, PositionwiseFeedForward
|
|
|
|
|
|
class FFTBlock(torch.nn.Module):
|
|
"""FFT Block"""
|
|
|
|
def __init__(self, d_model, n_head, d_k, d_v, d_inner, kernel_size, dropout=0.1):
|
|
super(FFTBlock, self).__init__()
|
|
self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
|
|
self.pos_ffn = PositionwiseFeedForward(
|
|
d_model, d_inner, kernel_size, dropout=dropout
|
|
)
|
|
|
|
def forward(self, enc_input, mask=None, slf_attn_mask=None):
|
|
enc_output, enc_slf_attn = self.slf_attn(
|
|
enc_input, enc_input, enc_input, mask=slf_attn_mask
|
|
)
|
|
enc_output = enc_output.masked_fill(mask.unsqueeze(-1), 0)
|
|
|
|
enc_output = self.pos_ffn(enc_output)
|
|
enc_output = enc_output.masked_fill(mask.unsqueeze(-1), 0)
|
|
|
|
return enc_output, enc_slf_attn
|
|
|
|
|
|
class ConvNorm(torch.nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=1,
|
|
stride=1,
|
|
padding=None,
|
|
dilation=1,
|
|
bias=True,
|
|
w_init_gain="linear",
|
|
):
|
|
super(ConvNorm, self).__init__()
|
|
|
|
if padding is None:
|
|
assert kernel_size % 2 == 1
|
|
padding = int(dilation * (kernel_size - 1) / 2)
|
|
|
|
self.conv = torch.nn.Conv1d(
|
|
in_channels,
|
|
out_channels,
|
|
kernel_size=kernel_size,
|
|
stride=stride,
|
|
padding=padding,
|
|
dilation=dilation,
|
|
bias=bias,
|
|
)
|
|
|
|
def forward(self, signal):
|
|
conv_signal = self.conv(signal)
|
|
|
|
return conv_signal
|
|
|
|
|
|
class PostNet(nn.Module):
|
|
"""
|
|
PostNet: Five 1-d convolution with 512 channels and kernel size 5
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
n_mel_channels=80,
|
|
postnet_embedding_dim=512,
|
|
postnet_kernel_size=5,
|
|
postnet_n_convolutions=5,
|
|
):
|
|
super(PostNet, self).__init__()
|
|
self.convolutions = nn.ModuleList()
|
|
|
|
self.convolutions.append(
|
|
nn.Sequential(
|
|
ConvNorm(
|
|
n_mel_channels,
|
|
postnet_embedding_dim,
|
|
kernel_size=postnet_kernel_size,
|
|
stride=1,
|
|
padding=int((postnet_kernel_size - 1) / 2),
|
|
dilation=1,
|
|
w_init_gain="tanh",
|
|
),
|
|
nn.BatchNorm1d(postnet_embedding_dim),
|
|
)
|
|
)
|
|
|
|
for i in range(1, postnet_n_convolutions - 1):
|
|
self.convolutions.append(
|
|
nn.Sequential(
|
|
ConvNorm(
|
|
postnet_embedding_dim,
|
|
postnet_embedding_dim,
|
|
kernel_size=postnet_kernel_size,
|
|
stride=1,
|
|
padding=int((postnet_kernel_size - 1) / 2),
|
|
dilation=1,
|
|
w_init_gain="tanh",
|
|
),
|
|
nn.BatchNorm1d(postnet_embedding_dim),
|
|
)
|
|
)
|
|
|
|
self.convolutions.append(
|
|
nn.Sequential(
|
|
ConvNorm(
|
|
postnet_embedding_dim,
|
|
n_mel_channels,
|
|
kernel_size=postnet_kernel_size,
|
|
stride=1,
|
|
padding=int((postnet_kernel_size - 1) / 2),
|
|
dilation=1,
|
|
w_init_gain="linear",
|
|
),
|
|
nn.BatchNorm1d(n_mel_channels),
|
|
)
|
|
)
|
|
|
|
def forward(self, x):
|
|
x = x.contiguous().transpose(1, 2)
|
|
|
|
for i in range(len(self.convolutions) - 1):
|
|
x = F.dropout(torch.tanh(self.convolutions[i](x)), 0.5, self.training)
|
|
x = F.dropout(self.convolutions[-1](x), 0.5, self.training)
|
|
|
|
x = x.contiguous().transpose(1, 2)
|
|
return x
|