Spaces:
Running
Running
| import torch | |
| from torch.autograd import Variable | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from nn_layers import linear_module, location_layer | |
| from utils import get_mask_from_lengths | |
| torch.manual_seed(1234) | |
| class AttentionNet(nn.Module): | |
| # 1024, 512, 128, 32, 31 | |
| def __init__(self, attention_rnn_dim, embedding_dim, attention_dim, | |
| attention_location_n_filters, attention_location_kernel_size): | |
| super(AttentionNet, self).__init__() | |
| self.query_layer = linear_module(attention_rnn_dim, attention_dim, | |
| bias=False, w_init_gain='tanh') | |
| # Projecting inputs into 128-D hidden representation | |
| self.memory_layer = linear_module(embedding_dim, attention_dim, bias=False, | |
| w_init_gain='tanh') | |
| # Projecting into 1-D scalar value | |
| self.v = linear_module(attention_dim, 1, bias=False) | |
| # Convolutional layers to obtain location features and projecting them into 128-D hidden representation | |
| self.location_layer = location_layer(attention_location_n_filters, | |
| attention_location_kernel_size, | |
| attention_dim) | |
| self.score_mask_value = -float("inf") | |
| def get_alignment_energies(self, query, processed_memory, | |
| attention_weights_cat): | |
| """ | |
| PARAMS | |
| ------ | |
| query: decoder output (batch, n_mel_channels * n_frames_per_step) | |
| processed_memory: processed encoder outputs (B, T_in, attention_dim) | |
| attention_weights_cat: cumulative and prev. att weights (B, 2, max_time) | |
| RETURNS | |
| ------- | |
| alignment (batch, max_time) | |
| """ | |
| processed_query = self.query_layer(query.unsqueeze(1)) | |
| processed_attention_weights = self.location_layer(attention_weights_cat) | |
| energies = self.v(torch.tanh( | |
| processed_query + processed_attention_weights + processed_memory)) | |
| energies = energies.squeeze(-1) # eliminates the third dimension of the tensor, which is 1. | |
| return energies | |
| def forward(self, attention_hidden_state, memory, processed_memory, | |
| attention_weights_cat, mask): | |
| """ | |
| PARAMS | |
| ------ | |
| attention_hidden_state: attention rnn last output | |
| memory: encoder outputs | |
| processed_memory: processed encoder outputs | |
| attention_weights_cat: previous and cummulative attention weights | |
| mask: binary mask for padded data | |
| """ | |
| alignment = self.get_alignment_energies( | |
| attention_hidden_state, processed_memory, attention_weights_cat) | |
| if mask is not None: | |
| alignment.data.masked_fill_(mask, self.score_mask_value) | |
| attention_weights = F.softmax(alignment, dim=1) | |
| # I think attention_weights is a [BxNUMENCINPUTS] so with unsequeeze(1): [Bx1xNUMENCINPUTS] and memory is | |
| # [BxNUMENCINPUTSx512] | |
| attention_context = torch.bmm(attention_weights.unsqueeze(1), memory) | |
| attention_context = attention_context.squeeze(1) | |
| return attention_context, attention_weights | |
| class Prenet(nn.Module): | |
| def __init__(self, in_dim, sizes): | |
| super(Prenet, self).__init__() | |
| in_sizes = [in_dim] + sizes[:-1] # all list values but the last one. The result is a list of the in_dim element | |
| # concatenated with sizes of layers (i.e. [80, 256]) | |
| self.layers = nn.ModuleList( | |
| [linear_module(in_size, out_size, bias=False) | |
| for (in_size, out_size) in zip(in_sizes, sizes)]) | |
| def forward(self, x): | |
| for linear in self.layers: | |
| x = F.dropout(F.relu(linear(x)), p=0.5, training=True) | |
| return x | |
| class Decoder(nn.Module): | |
| def __init__(self, tacotron_hyperparams): | |
| super(Decoder, self).__init__() | |
| self.n_mel_channels = tacotron_hyperparams['n_mel_channels'] | |
| self.n_frames_per_step = tacotron_hyperparams['number_frames_step'] | |
| self.encoder_embedding_dim = tacotron_hyperparams['encoder_embedding_dim'] | |
| self.attention_rnn_dim = tacotron_hyperparams['attention_rnn_dim'] # 1024 | |
| self.decoder_rnn_dim = tacotron_hyperparams['decoder_rnn_dim'] # 1024 | |
| self.prenet_dim = tacotron_hyperparams['prenet_dim'] | |
| self.max_decoder_steps = tacotron_hyperparams['max_decoder_steps'] | |
| # The threshold to decide whether stop or not stop decoding? | |
| self.gate_threshold = tacotron_hyperparams['gate_threshold'] | |
| self.p_attention_dropout = tacotron_hyperparams['p_attention_dropout'] | |
| self.p_decoder_dropout = tacotron_hyperparams['p_decoder_dropout'] | |
| # Define the prenet: there is only one frame per step, so input dim is the number of mel channels. | |
| # There are two fully connected layers: | |
| self.prenet = Prenet( | |
| tacotron_hyperparams['n_mel_channels'] * tacotron_hyperparams['number_frames_step'], | |
| [tacotron_hyperparams['prenet_dim'], tacotron_hyperparams['prenet_dim']]) | |
| # input_size: 1024 + 512 (output of first LSTM cell + attention_context) / hidden_size: 1024 | |
| self.attention_rnn = nn.LSTMCell( | |
| tacotron_hyperparams['prenet_dim'] + tacotron_hyperparams['encoder_embedding_dim'], | |
| tacotron_hyperparams['attention_rnn_dim']) | |
| # return attention_weights and attention_context. Does the alignments. | |
| self.attention_layer = AttentionNet( | |
| tacotron_hyperparams['attention_rnn_dim'], tacotron_hyperparams['encoder_embedding_dim'], | |
| tacotron_hyperparams['attention_dim'], tacotron_hyperparams['attention_location_n_filters'], | |
| tacotron_hyperparams['attention_location_kernel_size']) | |
| # input_size: 256 + 512 (attention_context + prenet_info), hidden_size: 1024 | |
| self.decoder_rnn = nn.LSTMCell( | |
| tacotron_hyperparams['attention_rnn_dim'] + tacotron_hyperparams['encoder_embedding_dim'], | |
| tacotron_hyperparams['decoder_rnn_dim'], 1) | |
| # (LSTM output)1024 + (attention_context)512, out_dim: number of mel channels. Last linear projection that | |
| # generates an output decoder spectral frame. | |
| self.linear_projection = linear_module( | |
| tacotron_hyperparams['decoder_rnn_dim'] + tacotron_hyperparams['encoder_embedding_dim'], | |
| tacotron_hyperparams['n_mel_channels']*tacotron_hyperparams['number_frames_step']) | |
| # decision whether to continue decoding. | |
| self.gate_layer = linear_module( | |
| tacotron_hyperparams['decoder_rnn_dim'] + tacotron_hyperparams['encoder_embedding_dim'], 1, | |
| bias=True, w_init_gain='sigmoid') | |
| def get_go_frame(self, memory): | |
| """ Gets all zeros frames to use as first decoder input | |
| PARAMS | |
| ------ | |
| memory: decoder outputs | |
| RETURNS | |
| ------- | |
| decoder_input: all zeros frames | |
| """ | |
| B = memory.size(0) | |
| decoder_input = Variable(memory.data.new( | |
| B, self.n_mel_channels * self.n_frames_per_step).zero_()) | |
| return decoder_input | |
| def initialize_decoder_states(self, memory, mask): | |
| """ Initializes attention rnn states, decoder rnn states, attention | |
| weights, attention cumulative weights, attention context, stores memory | |
| and stores processed memory | |
| PARAMS | |
| ------ | |
| memory: Encoder outputs | |
| mask: Mask for padded data if training, expects None for inference | |
| """ | |
| B = memory.size(0) | |
| MAX_TIME = memory.size(1) | |
| self.attention_hidden = Variable(memory.data.new( | |
| B, self.attention_rnn_dim).zero_()) | |
| self.attention_cell = Variable(memory.data.new( | |
| B, self.attention_rnn_dim).zero_()) | |
| self.decoder_hidden = Variable(memory.data.new( | |
| B, self.decoder_rnn_dim).zero_()) | |
| self.decoder_cell = Variable(memory.data.new( | |
| B, self.decoder_rnn_dim).zero_()) | |
| self.attention_weights = Variable(memory.data.new( | |
| B, MAX_TIME).zero_()) | |
| self.attention_weights_cum = Variable(memory.data.new( | |
| B, MAX_TIME).zero_()) | |
| self.attention_context = Variable(memory.data.new( | |
| B, self.encoder_embedding_dim).zero_()) | |
| self.memory = memory | |
| self.processed_memory = self.attention_layer.memory_layer(memory) | |
| self.mask = mask | |
| def parse_decoder_inputs(self, decoder_inputs): | |
| """ Prepares decoder inputs, i.e. mel outputs | |
| PARAMS | |
| ------ | |
| decoder_inputs: inputs used for teacher-forced training, i.e. mel-specs | |
| RETURNS | |
| ------- | |
| inputs: processed decoder inputs | |
| """ | |
| # (B, n_mel_channels, T_out) -> (B, T_out, n_mel_channels) | |
| decoder_inputs = decoder_inputs.transpose(1, 2) | |
| # reshape decoder inputs in case we want to work with more than 1 frame per step (chunks). Otherwise, this next | |
| # line does not just do anything | |
| decoder_inputs = decoder_inputs.view( | |
| decoder_inputs.size(0), | |
| int(decoder_inputs.size(1)/self.n_frames_per_step), -1) | |
| # (B, T_out, n_mel_channels) -> (T_out, B, n_mel_channels) | |
| decoder_inputs = decoder_inputs.transpose(0, 1) | |
| return decoder_inputs | |
| def parse_decoder_outputs(self, mel_outputs, gate_outputs, alignments): | |
| """ Prepares decoder outputs for output | |
| PARAMS | |
| ------ | |
| mel_outputs: | |
| gate_outputs: gate output energies | |
| alignments: | |
| RETURNS | |
| ------- | |
| mel_outputs: | |
| gate_outpust: gate output energies | |
| alignments: | |
| """ | |
| # (T_out, B) -> (B, T_out) | |
| alignments = torch.stack(alignments).transpose(0, 1) | |
| # (T_out, B) -> (B, T_out) | |
| gate_outputs = torch.stack(gate_outputs).transpose(0, 1) | |
| gate_outputs = gate_outputs.contiguous() | |
| # (T_out, B, n_mel_channels) -> (B, T_out, n_mel_channels) | |
| mel_outputs = torch.stack(mel_outputs).transpose(0, 1).contiguous() | |
| # decouple frames per step | |
| mel_outputs = mel_outputs.view( | |
| mel_outputs.size(0), -1, self.n_mel_channels) | |
| # (B, T_out, n_mel_channels) -> (B, n_mel_channels, T_out) | |
| mel_outputs = mel_outputs.transpose(1, 2) | |
| return mel_outputs, gate_outputs, alignments | |
| def decode(self, decoder_input): | |
| """ Decoder step using stored states, attention and memory | |
| PARAMS | |
| ------ | |
| decoder_input: previous mel output | |
| RETURNS | |
| ------- | |
| mel_output: | |
| gate_output: gate output energies | |
| attention_weights: | |
| """ | |
| # concatenates [Bx1024] and [Bx512]. All dimensions match except 1 (torch.cat -1) | |
| # concatenate the i-th decoder hidden state together with the i-th attention context | |
| cell_input = torch.cat((decoder_input, self.attention_context), -1) | |
| # the previous input is for the following LSTM cell, initialized with zeroes the hidden states and the cell | |
| # state. | |
| # compute the (i+1)th attention hidden state based on the i-th decoder hidden state and attention context. | |
| self.attention_hidden, self.attention_cell = self.attention_rnn( | |
| cell_input, (self.attention_hidden, self.attention_cell)) | |
| self.attention_hidden = F.dropout(self.attention_hidden, self.p_attention_dropout, self.training) | |
| self.attention_cell = F.dropout(self.attention_cell, self.p_attention_dropout, self.training) | |
| # concatenate the i-th state attention weights together with the cumulated from previous states to compute | |
| # (i+1)th state | |
| attention_weights_cat = torch.cat( | |
| (self.attention_weights.unsqueeze(1), | |
| self.attention_weights_cum.unsqueeze(1)), dim=1) | |
| # compute (i+1)th attention context and provide (i+1)th attention weights based on the (i+1)th attention hidden | |
| # state and (i)th and prev. weights | |
| self.attention_context, self.attention_weights = self.attention_layer( | |
| self.attention_hidden, self.memory, self.processed_memory, | |
| attention_weights_cat, self.mask) | |
| # cumulate attention_weights adding the (i+1)th to compute (i+2)th state | |
| self.attention_weights_cum += self.attention_weights | |
| decoder_input = torch.cat((self.attention_hidden, self.attention_context), -1) | |
| self.decoder_hidden, self.decoder_cell = self.decoder_rnn(decoder_input, | |
| (self.decoder_hidden, self.decoder_cell)) | |
| self.decoder_hidden = F.dropout(self.decoder_hidden, self.p_decoder_dropout, self.training) | |
| self.decoder_cell = F.dropout(self.decoder_cell, self.p_decoder_dropout, self.training) | |
| decoder_hidden_attention_context = torch.cat((self.decoder_hidden, self.attention_context), dim=1) | |
| decoder_output = self.linear_projection(decoder_hidden_attention_context) | |
| gate_prediction = self.gate_layer(decoder_hidden_attention_context) | |
| return decoder_output, gate_prediction, self.attention_weights | |
| """ | |
| # the decoder_output from ith step passes through the pre-net to compute new decoder hidden state and attention_ | |
| # context (i+1)th | |
| prenet_output = self.prenet(decoder_input) | |
| # the decoder_input now is the concatenation of the pre-net output and the new (i+1)th attention_context | |
| decoder_input = torch.cat((prenet_output, self.attention_context), -1) | |
| # another LSTM Cell to compute the decoder hidden (i+1)th state from the decoder_input | |
| self.decoder_hidden, self.decoder_cell = self.decoder_rnn( | |
| decoder_input, (self.decoder_hidden, self.decoder_cell)) | |
| # with new attention_context we concatenate again with the new (i+1)th decoder_hidden state. | |
| decoder_hidden_attention_context = torch.cat( | |
| (self.decoder_hidden, self.attention_context), dim=1) | |
| # the (i+1)th output is a linear projection of the decoder hidden state with a weight matrix plus bias. | |
| decoder_output = self.linear_projection( | |
| decoder_hidden_attention_context) | |
| # check whether (i+1)th state is the last of the sequence | |
| gate_prediction = self.gate_layer(decoder_hidden_attention_context) | |
| return decoder_output, gate_prediction, self.attention_weights""" | |
| def forward(self, memory, decoder_inputs, memory_lengths): | |
| """ Decoder forward pass for training | |
| PARAMS | |
| ------ | |
| memory: Encoder outputs | |
| decoder_inputs: Decoder inputs for teacher forcing. i.e. mel-specs | |
| memory_lengths: Encoder output lengths for attention masking. | |
| RETURNS | |
| ------- | |
| mel_outputs: mel outputs from the decoder | |
| gate_outputs: gate outputs from the decoder | |
| alignments: sequence of attention weights from the decoder | |
| """ | |
| decoder_input = self.get_go_frame(memory).unsqueeze(0) | |
| decoder_inputs = self.parse_decoder_inputs(decoder_inputs) | |
| decoder_inputs = torch.cat((decoder_input, decoder_inputs), dim=0) | |
| decoder_inputs = self.prenet(decoder_inputs) | |
| self.initialize_decoder_states( | |
| memory, mask=~get_mask_from_lengths(memory_lengths)) | |
| mel_outputs, gate_outputs, alignments = [], [], [] | |
| while len(mel_outputs) < decoder_inputs.size(0) - 1: | |
| decoder_input = decoder_inputs[len(mel_outputs)] | |
| mel_output, gate_output, attention_weights = self.decode( | |
| decoder_input) | |
| # a class list, when += means concatenation of vectors | |
| mel_outputs += [mel_output.squeeze(1)] | |
| gate_outputs += [gate_output.squeeze()] | |
| alignments += [attention_weights] | |
| # getting the frame indexing from reference mel frames to pass it as the new input of the next decoding | |
| # step: Teacher Forcing! | |
| # Takes each time_step of sequences of all mini-batch samples (i.e. [48, 80] as the decoder_inputs is | |
| # parsed as [189, 48, 80]). | |
| mel_outputs, gate_outputs, alignments = self.parse_decoder_outputs( | |
| mel_outputs, gate_outputs, alignments) | |
| return mel_outputs, gate_outputs, alignments | |
| def inference(self, memory): | |
| """ Decoder inference | |
| PARAMS | |
| ------ | |
| memory: Encoder outputs | |
| RETURNS | |
| ------- | |
| mel_outputs: mel outputs from the decoder | |
| gate_outputs: gate outputs from the decoder | |
| alignments: sequence of attention weights from the decoder | |
| """ | |
| decoder_input = self.get_go_frame(memory) | |
| self.initialize_decoder_states(memory, mask=None) | |
| mel_outputs, gate_outputs, alignments = [], [], [] | |
| while True: | |
| decoder_input = self.prenet(decoder_input) | |
| mel_output, gate_output, alignment = self.decode(decoder_input) | |
| mel_outputs += [mel_output.squeeze(1)] | |
| gate_outputs += [gate_output] | |
| alignments += [alignment] | |
| if torch.sigmoid(gate_output.data) > self.gate_threshold: | |
| break | |
| elif len(mel_outputs) == self.max_decoder_steps: | |
| print("Warning! Reached max decoder steps") | |
| break | |
| decoder_input = mel_output | |
| mel_outputs, gate_outputs, alignments = self.parse_decoder_outputs( | |
| mel_outputs, gate_outputs, alignments) | |
| return mel_outputs, gate_outputs, alignments | |