mirror of https://github.com/coqui-ai/TTS.git
335 lines
12 KiB
Python
335 lines
12 KiB
Python
# coding: utf-8
|
|
import torch
|
|
from torch.autograd import Variable
|
|
from torch import nn
|
|
|
|
from .attention import AttentionRNN
|
|
from .attention import get_mask_from_lengths
|
|
|
|
|
|
class Prenet(nn.Module):
|
|
r""" Prenet as explained at https://arxiv.org/abs/1703.10135.
|
|
It creates as many layers as given by 'out_features'
|
|
|
|
Args:
|
|
in_features (int): size of the input vector
|
|
out_features (int or list): size of each output sample.
|
|
If it is a list, for each value, there is created a new layer.
|
|
"""
|
|
|
|
def __init__(self, in_features, out_features=[256, 128]):
|
|
super(Prenet, self).__init__()
|
|
in_features = [in_features] + out_features[:-1]
|
|
self.layers = nn.ModuleList(
|
|
[nn.Linear(in_size, out_size)
|
|
for (in_size, out_size) in zip(in_features, out_features)])
|
|
self.relu = nn.ReLU()
|
|
self.dropout = nn.Dropout(0.5)
|
|
|
|
def forward(self, inputs):
|
|
for linear in self.layers:
|
|
inputs = self.dropout(self.relu(linear(inputs)))
|
|
return inputs
|
|
|
|
|
|
class BatchNormConv1d(nn.Module):
|
|
r"""A wrapper for Conv1d with BatchNorm. It sets the activation
|
|
function between Conv and BatchNorm layers. BatchNorm layer
|
|
is initialized with the TF default values for momentum and eps.
|
|
|
|
Args:
|
|
in_channels: size of each input sample
|
|
out_channels: size of each output samples
|
|
kernel_size: kernel size of conv filters
|
|
stride: stride of conv filters
|
|
padding: padding of conv filters
|
|
activation: activation function set b/w Conv1d and BatchNorm
|
|
|
|
Shapes:
|
|
- input: batch x dims
|
|
- output: batch x dims
|
|
"""
|
|
|
|
def __init__(self, in_channels, out_channels, kernel_size, stride, padding,
|
|
activation=None):
|
|
super(BatchNormConv1d, self).__init__()
|
|
self.conv1d = nn.Conv1d(in_channels, out_channels,
|
|
kernel_size=kernel_size,
|
|
stride=stride, padding=padding, bias=False)
|
|
# Following tensorflow's default parameters
|
|
self.bn = nn.BatchNorm1d(out_channels, momentum=0.99, eps=1e-3)
|
|
self.activation = activation
|
|
|
|
def forward(self, x):
|
|
x = self.conv1d(x)
|
|
if self.activation is not None:
|
|
x = self.activation(x)
|
|
return self.bn(x)
|
|
|
|
|
|
class Highway(nn.Module):
|
|
def __init__(self, in_size, out_size):
|
|
super(Highway, self).__init__()
|
|
self.H = nn.Linear(in_size, out_size)
|
|
self.H.bias.data.zero_()
|
|
self.T = nn.Linear(in_size, out_size)
|
|
self.T.bias.data.fill_(-1)
|
|
self.relu = nn.ReLU()
|
|
self.sigmoid = nn.Sigmoid()
|
|
|
|
def forward(self, inputs):
|
|
H = self.relu(self.H(inputs))
|
|
T = self.sigmoid(self.T(inputs))
|
|
return H * T + inputs * (1.0 - T)
|
|
|
|
|
|
class CBHG(nn.Module):
|
|
"""CBHG module: a recurrent neural network composed of:
|
|
- 1-d convolution banks
|
|
- Highway networks + residual connections
|
|
- Bidirectional gated recurrent units
|
|
|
|
Args:
|
|
in_features (int): sample size
|
|
K (int): max filter size in conv bank
|
|
projections (list): conv channel sizes for conv projections
|
|
num_highways (int): number of highways layers
|
|
|
|
Shapes:
|
|
- input: batch x time x dim
|
|
- output: batch x time x dim*2
|
|
"""
|
|
|
|
def __init__(self, in_features, K=16, projections=[128, 128], num_highways=4):
|
|
super(CBHG, self).__init__()
|
|
self.in_features = in_features
|
|
self.relu = nn.ReLU()
|
|
|
|
# list of conv1d bank with filter size k=1...K
|
|
# TODO: try dilational layers instead
|
|
self.conv1d_banks = nn.ModuleList(
|
|
[BatchNormConv1d(in_features, in_features, kernel_size=k, stride=1,
|
|
padding=k // 2, activation=self.relu)
|
|
for k in range(1, K + 1)])
|
|
|
|
# max pooling of conv bank
|
|
# TODO: try average pooling OR larger kernel size
|
|
self.max_pool1d = nn.MaxPool1d(kernel_size=2, stride=1, padding=1)
|
|
|
|
out_features = [K * in_features] + projections[:-1]
|
|
activations = [self.relu] * (len(projections) - 1)
|
|
activations += [None]
|
|
|
|
# setup conv1d projection layers
|
|
layer_set = []
|
|
for (in_size, out_size, ac) in zip(out_features, projections, activations):
|
|
layer = BatchNormConv1d(in_size, out_size, kernel_size=3, stride=1,
|
|
padding=1, activation=ac)
|
|
layer_set.append(layer)
|
|
self.conv1d_projections = nn.ModuleList(layer_set)
|
|
|
|
# setup Highway layers
|
|
self.pre_highway = nn.Linear(projections[-1], in_features, bias=False)
|
|
self.highways = nn.ModuleList(
|
|
[Highway(in_features, in_features) for _ in range(num_highways)])
|
|
|
|
# bi-directional GPU layer
|
|
self.gru = nn.GRU(
|
|
in_features, in_features, 1, batch_first=True, bidirectional=True)
|
|
|
|
def forward(self, inputs):
|
|
# (B, T_in, in_features)
|
|
x = inputs
|
|
|
|
# Needed to perform conv1d on time-axis
|
|
# (B, in_features, T_in)
|
|
if x.size(-1) == self.in_features:
|
|
x = x.transpose(1, 2)
|
|
|
|
T = x.size(-1)
|
|
|
|
# (B, in_features*K, T_in)
|
|
# Concat conv1d bank outputs
|
|
outs = []
|
|
for conv1d in self.conv1d_banks:
|
|
out = conv1d(x)
|
|
out = out[:, :, :T]
|
|
outs.append(out)
|
|
|
|
x = torch.cat(outs, dim=1)
|
|
assert x.size(1) == self.in_features * len(self.conv1d_banks)
|
|
|
|
x = self.max_pool1d(x)[:, :, :T]
|
|
|
|
for conv1d in self.conv1d_projections:
|
|
x = conv1d(x)
|
|
|
|
# (B, T_in, in_features)
|
|
# Back to the original shape
|
|
x = x.transpose(1, 2)
|
|
|
|
if x.size(-1) != self.in_features:
|
|
x = self.pre_highway(x)
|
|
|
|
# Residual connection
|
|
# TODO: try residual scaling as in Deep Voice 3
|
|
# TODO: try plain residual layers
|
|
x += inputs
|
|
for highway in self.highways:
|
|
x = highway(x)
|
|
|
|
# (B, T_in, in_features*2)
|
|
# TODO: replace GRU with convolution as in Deep Voice 3
|
|
self.gru.flatten_parameters()
|
|
outputs, _ = self.gru(x)
|
|
return outputs
|
|
|
|
|
|
class Encoder(nn.Module):
|
|
r"""Encapsulate Prenet and CBHG modules for encoder"""
|
|
|
|
def __init__(self, in_features):
|
|
super(Encoder, self).__init__()
|
|
self.prenet = Prenet(in_features, out_features=[256, 128])
|
|
self.cbhg = CBHG(128, K=16, projections=[128, 128])
|
|
|
|
def forward(self, inputs):
|
|
r"""
|
|
Args:
|
|
inputs (FloatTensor): embedding features
|
|
|
|
Shapes:
|
|
- inputs: batch x time x in_features
|
|
- outputs: batch x time x 128*2
|
|
"""
|
|
inputs = self.prenet(inputs)
|
|
return self.cbhg(inputs)
|
|
|
|
|
|
class Decoder(nn.Module):
|
|
r"""Decoder module.
|
|
|
|
Args:
|
|
in_features (int): input vector (encoder output) sample size.
|
|
memory_dim (int): memory vector (prev. time-step output) sample size.
|
|
r (int): number of outputs per time step.
|
|
eps (float): threshold for detecting the end of a sentence.
|
|
"""
|
|
|
|
def __init__(self, in_features, memory_dim, r, eps=0.05, mode='train'):
|
|
super(Decoder, self).__init__()
|
|
self.mode = mode
|
|
self.max_decoder_steps = 200
|
|
self.memory_dim = memory_dim
|
|
self.eps = eps
|
|
self.r = r
|
|
# memory -> |Prenet| -> processed_memory
|
|
self.prenet = Prenet(memory_dim * r, out_features=[256, 128])
|
|
# processed_inputs, processed_memory -> |Attention| -> Attention, Alignment, RNN_State
|
|
self.attention_rnn = AttentionRNN(256, in_features, 128)
|
|
# (processed_memory | attention context) -> |Linear| -> decoder_RNN_input
|
|
self.project_to_decoder_in = nn.Linear(256+in_features, 256)
|
|
# decoder_RNN_input -> |RNN| -> RNN_state
|
|
self.decoder_rnns = nn.ModuleList(
|
|
[nn.GRUCell(256, 256) for _ in range(2)])
|
|
# RNN_state -> |Linear| -> mel_spec
|
|
self.proj_to_mel = nn.Linear(256, memory_dim * r)
|
|
|
|
def forward(self, inputs, memory=None):
|
|
"""
|
|
Decoder forward step.
|
|
|
|
If decoder inputs are not given (e.g., at testing time), as noted in
|
|
Tacotron paper, greedy decoding is adapted.
|
|
|
|
Args:
|
|
inputs: Encoder outputs.
|
|
memory (None): Decoder memory (autoregression. If None (at eval-time),
|
|
decoder outputs are used as decoder inputs. If None, it uses the last
|
|
output as the input.
|
|
|
|
Shapes:
|
|
- inputs: batch x time x encoder_out_dim
|
|
- memory: batch x #mels_pecs x mel_spec_dim
|
|
"""
|
|
B = inputs.size(0)
|
|
# Run greedy decoding if memory is None
|
|
greedy = not self.training
|
|
if memory is not None:
|
|
# Grouping multiple frames if necessary
|
|
if memory.size(-1) == self.memory_dim:
|
|
memory = memory.view(B, memory.size(1) // self.r, -1)
|
|
" !! Dimension mismatch {} vs {} * {}".format(memory.size(-1),
|
|
self.memory_dim, self.r)
|
|
T_decoder = memory.size(1)
|
|
# go frame - 0 frames tarting the sequence
|
|
initial_memory = Variable(
|
|
inputs.data.new(B, self.memory_dim * self.r).zero_())
|
|
# Init decoder states
|
|
attention_rnn_hidden = Variable(
|
|
inputs.data.new(B, 256).zero_())
|
|
decoder_rnn_hiddens = [Variable(
|
|
inputs.data.new(B, 256).zero_())
|
|
for _ in range(len(self.decoder_rnns))]
|
|
current_context_vec = Variable(
|
|
inputs.data.new(B, 256).zero_())
|
|
# Time first (T_decoder, B, memory_dim)
|
|
if memory is not None:
|
|
memory = memory.transpose(0, 1)
|
|
outputs = []
|
|
alignments = []
|
|
t = 0
|
|
memory_input = initial_memory
|
|
while True:
|
|
if t > 0:
|
|
if greedy:
|
|
memory_input = outputs[-1]
|
|
else:
|
|
# TODO: try sampled teacher forcing
|
|
# combine prev. model output and prev. real target
|
|
# memory_input = torch.div(outputs[-1] + memory[t-1], 2.0)
|
|
# add a random noise
|
|
# noise = torch.autograd.Variable(
|
|
# memory_input.data.new(memory_input.size()).normal_(0.0, 0.5))
|
|
# memory_input = memory_input + noise
|
|
memory_input = memory[t-1]
|
|
# Prenet
|
|
processed_memory = self.prenet(memory_input)
|
|
# Attention RNN
|
|
attention_rnn_hidden, current_context_vec, alignment = self.attention_rnn(
|
|
processed_memory, current_context_vec, attention_rnn_hidden, inputs)
|
|
# Concat RNN output and attention context vector
|
|
decoder_input = self.project_to_decoder_in(
|
|
torch.cat((attention_rnn_hidden, current_context_vec), -1))
|
|
# Pass through the decoder RNNs
|
|
for idx in range(len(self.decoder_rnns)):
|
|
decoder_rnn_hiddens[idx] = self.decoder_rnns[idx](
|
|
decoder_input, decoder_rnn_hiddens[idx])
|
|
# Residual connectinon
|
|
decoder_input = decoder_rnn_hiddens[idx] + decoder_input
|
|
output = decoder_input
|
|
# predict mel vectors from decoder vectors
|
|
output = self.proj_to_mel(output)
|
|
outputs += [output]
|
|
alignments += [alignment]
|
|
t += 1
|
|
if (not greedy and self.training) or (greedy and memory is not None):
|
|
if t >= T_decoder:
|
|
break
|
|
else:
|
|
if t > 1 and is_end_of_frames(output, self.eps):
|
|
break
|
|
elif t > self.max_decoder_steps:
|
|
print(" !! Decoder stopped with 'max_decoder_steps'. \
|
|
Something is probably wrong.")
|
|
break
|
|
assert greedy or len(outputs) == T_decoder
|
|
# Back to batch first
|
|
alignments = torch.stack(alignments).transpose(0, 1)
|
|
outputs = torch.stack(outputs).transpose(0, 1).contiguous()
|
|
return outputs, alignments
|
|
|
|
|
|
def is_end_of_frames(output, eps=0.2): # 0.2
|
|
return (output.data <= eps).all()
|