import os import json import copy import math from collections import OrderedDict import torch import torch.nn as nn import numpy as np import torch.nn.functional as F from torch.autograd import Variable from utils.tools import get_mask_from_lengths, pad from .blocks import ( LinearNorm, ConvNorm, ) from text.symbols import symbols device = torch.device("cuda" if torch.cuda.is_available() else "cpu") from typing import Optional, Tuple from numba import jit, prange LRELU_SLOPE = 0.1 def sequence_mask(length, max_length=None): if max_length is None: max_length = length.max() x = torch.arange(max_length, dtype=length.dtype, device=length.device) return x.unsqueeze(0) < length.unsqueeze(1) @jit(nopython=True) def mas_width1(attn_map): """mas with hardcoded width=1""" # assumes mel x text opt = np.zeros_like(attn_map) attn_map = np.log(attn_map) attn_map[0, 1:] = -np.inf log_p = np.zeros_like(attn_map) log_p[0, :] = attn_map[0, :] prev_ind = np.zeros_like(attn_map, dtype=np.int64) for i in range(1, attn_map.shape[0]): for j in range(attn_map.shape[1]): # for each text dim prev_log = log_p[i - 1, j] prev_j = j if j - 1 >= 0 and log_p[i - 1, j - 1] >= log_p[i - 1, j]: prev_log = log_p[i - 1, j - 1] prev_j = j - 1 log_p[i, j] = attn_map[i, j] + prev_log prev_ind[i, j] = prev_j # now backtrack curr_text_idx = attn_map.shape[1] - 1 for i in range(attn_map.shape[0] - 1, -1, -1): opt[i, curr_text_idx] = 1 curr_text_idx = prev_ind[i, curr_text_idx] opt[0, curr_text_idx] = 1 return opt @jit(nopython=True) def b_mas(b_attn_map, in_lens, out_lens, width=1): assert width == 1 attn_out = np.zeros_like(b_attn_map) for b in prange(b_attn_map.shape[0]): out = mas_width1(b_attn_map[b, 0, : out_lens[b], : in_lens[b]]) attn_out[b, 0, : out_lens[b], : in_lens[b]] = out return attn_out class PartialConv1d(torch.nn.Conv1d): """ Zero padding creates a unique identifier for where the edge of the data is, such that the model can almost always identify exactly where it is relative to either edge given a sufficient receptive field. Partial padding goes to some lengths to remove this affect. """ def __init__(self, *args, **kwargs): super(PartialConv1d, self).__init__(*args, **kwargs) weight_maskUpdater = torch.ones(1, 1, self.kernel_size[0]) self.register_buffer("weight_maskUpdater", weight_maskUpdater, persistent=False) slide_winsize = torch.tensor(self.weight_maskUpdater.shape[1] * self.weight_maskUpdater.shape[2]) self.register_buffer("slide_winsize", slide_winsize, persistent=False) if self.bias is not None: bias_view = self.bias.view(1, self.out_channels, 1) self.register_buffer('bias_view', bias_view, persistent=False) # caching part self.last_size = (-1, -1, -1) update_mask = torch.ones(1, 1, 1) self.register_buffer('update_mask', update_mask, persistent=False) mask_ratio = torch.ones(1, 1, 1) self.register_buffer('mask_ratio', mask_ratio, persistent=False) self.partial: bool = True def calculate_mask(self, input: torch.Tensor, mask_in: Optional[torch.Tensor]): with torch.no_grad(): if mask_in is None: mask = torch.ones(1, 1, input.shape[2], dtype=input.dtype, device=input.device) else: mask = mask_in update_mask = F.conv1d( mask, self.weight_maskUpdater, bias=None, stride=self.stride, padding=self.padding, dilation=self.dilation, groups=1, ) # for mixed precision training, change 1e-8 to 1e-6 mask_ratio = self.slide_winsize / (update_mask + 1e-6) update_mask = torch.clamp(update_mask, 0, 1) mask_ratio = torch.mul(mask_ratio.to(update_mask), update_mask) return torch.mul(input, mask), mask_ratio, update_mask def forward_aux(self, input: torch.Tensor, mask_ratio: torch.Tensor, update_mask: torch.Tensor) -> torch.Tensor: assert len(input.shape) == 3 raw_out = self._conv_forward(input, self.weight, self.bias) if self.bias is not None: output = torch.mul(raw_out - self.bias_view, mask_ratio) + self.bias_view output = torch.mul(output, update_mask) else: output = torch.mul(raw_out, mask_ratio) return output @torch.jit.ignore def forward_with_cache(self, input: torch.Tensor, mask_in: Optional[torch.Tensor] = None) -> torch.Tensor: use_cache = not (torch.jit.is_tracing() or torch.onnx.is_in_onnx_export()) cache_hit = use_cache and mask_in is None and self.last_size == input.shape if cache_hit: mask_ratio = self.mask_ratio update_mask = self.update_mask else: input, mask_ratio, update_mask = self.calculate_mask(input, mask_in) if use_cache: # if a mask is input, or tensor shape changed, update mask ratio self.last_size = tuple(input.shape) self.update_mask = update_mask self.mask_ratio = mask_ratio return self.forward_aux(input, mask_ratio, update_mask) def forward_no_cache(self, input: torch.Tensor, mask_in: Optional[torch.Tensor] = None) -> torch.Tensor: if self.partial: input, mask_ratio, update_mask = self.calculate_mask(input, mask_in) return self.forward_aux(input, mask_ratio, update_mask) else: if mask_in is not None: input = torch.mul(input, mask_in) return self._conv_forward(input, self.weight, self.bias) def forward(self, input: torch.Tensor, mask_in: Optional[torch.Tensor] = None) -> torch.Tensor: if self.partial: return self.forward_with_cache(input, mask_in) else: if mask_in is not None: input = torch.mul(input, mask_in) return self._conv_forward(input, self.weight, self.bias) 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', use_partial_padding: bool = False, use_weight_norm: bool = False, norm_fn=None, ): super(ConvNorm, self).__init__() if padding is None: assert kernel_size % 2 == 1 padding = int(dilation * (kernel_size - 1) / 2) self.use_partial_padding: bool = use_partial_padding conv = PartialConv1d( in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, ) conv.partial = use_partial_padding torch.nn.init.xavier_uniform_(conv.weight, gain=torch.nn.init.calculate_gain(w_init_gain)) if use_weight_norm: conv = torch.nn.utils.weight_norm(conv) if norm_fn is not None: self.norm = norm_fn(out_channels, affine=True) else: self.norm = None self.conv = conv def forward(self, input: torch.Tensor, mask_in: Optional[torch.Tensor] = None) -> torch.Tensor: ret = self.conv(input, mask_in) if self.norm is not None: ret = self.norm(ret) return ret def binarize_attention_parallel(attn, in_lens, out_lens): """For training purposes only. Binarizes attention with MAS. These will no longer receive a gradient. Args: attn: B x 1 x max_mel_len x max_text_len """ with torch.no_grad(): attn_cpu = attn.data.cpu().numpy() attn_out = b_mas(attn_cpu, in_lens.cpu().numpy(), out_lens.cpu().numpy(), width=1) return torch.from_numpy(attn_out).to(attn.device) class AlignmentEncoder(torch.nn.Module): """Module for alignment text and mel spectrogram. """ def __init__( self, n_mel_channels=80, n_text_channels=512, n_att_channels=80, temperature=0.0005, ): super().__init__() self.temperature = temperature self.softmax = torch.nn.Softmax(dim=3) self.log_softmax = torch.nn.LogSoftmax(dim=3) self.key_proj = nn.Sequential( ConvNorm(n_text_channels, n_text_channels * 2, kernel_size=3, bias=True, w_init_gain='relu'), torch.nn.ReLU(), ConvNorm(n_text_channels * 2, n_att_channels, kernel_size=1, bias=True), ) self.query_proj = nn.Sequential( ConvNorm(n_mel_channels, n_mel_channels * 2, kernel_size=3, bias=True, w_init_gain='relu'), torch.nn.ReLU(), ConvNorm(n_mel_channels * 2, n_mel_channels, kernel_size=1, bias=True), torch.nn.ReLU(), ConvNorm(n_mel_channels, n_att_channels, kernel_size=1, bias=True), ) def get_dist(self, keys, queries, mask=None): """Calculation of distance matrix. Args: queries (torch.tensor): B x C x T1 tensor (probably going to be mel data). keys (torch.tensor): B x C2 x T2 tensor (text data). mask (torch.tensor): B x T2 x 1 tensor, binary mask for variable length entries and also can be used for ignoring unnecessary elements from keys in the resulting distance matrix (True = mask element, False = leave unchanged). Output: dist (torch.tensor): B x T1 x T2 tensor. """ keys_enc = self.key_proj(keys) # B x n_attn_dims x T2 queries_enc = self.query_proj(queries) # B x n_attn_dims x T1 attn = (queries_enc[:, :, :, None] - keys_enc[:, :, None]) ** 2 # B x n_attn_dims x T1 x T2 dist = attn.sum(1, keepdim=True) # B x 1 x T1 x T2 if mask is not None: dist.data.masked_fill_(mask.permute(0, 2, 1).unsqueeze(2), float("inf")) return dist.squeeze(1) @staticmethod def get_durations(attn_soft, text_len, spect_len): """Calculation of durations. Args: attn_soft (torch.tensor): B x 1 x T1 x T2 tensor. text_len (torch.tensor): B tensor, lengths of text. spect_len (torch.tensor): B tensor, lengths of mel spectrogram. """ attn_hard = binarize_attention_parallel(attn_soft, text_len, spect_len) durations = attn_hard.sum(2)[:, 0, :] assert torch.all(torch.eq(durations.sum(dim=1), spect_len)) return durations @staticmethod def get_mean_dist_by_durations(dist, durations, mask=None): """Select elements from the distance matrix for the given durations and mask and return mean distance. Args: dist (torch.tensor): B x T1 x T2 tensor. durations (torch.tensor): B x T2 tensor. Dim T2 should sum to T1. mask (torch.tensor): B x T2 x 1 binary mask for variable length entries and also can be used for ignoring unnecessary elements in dist by T2 dim (True = mask element, False = leave unchanged). Output: mean_dist (torch.tensor): B x 1 tensor. """ batch_size, t1_size, t2_size = dist.size() assert torch.all(torch.eq(durations.sum(dim=1), t1_size)) if mask is not None: dist = dist.masked_fill(mask.permute(0, 2, 1).unsqueeze(2), 0) # TODO(oktai15): make it more efficient mean_dist_by_durations = [] for dist_idx in range(batch_size): mean_dist_by_durations.append( torch.mean( dist[ dist_idx, torch.arange(t1_size), torch.repeat_interleave(torch.arange(t2_size), repeats=durations[dist_idx]), ] ) ) return torch.tensor(mean_dist_by_durations, dtype=dist.dtype, device=dist.device) @staticmethod def get_mean_distance_for_word(l2_dists, durs, start_token, num_tokens): """Calculates the mean distance between text and audio embeddings given a range of text tokens. Args: l2_dists (torch.tensor): L2 distance matrix from Aligner inference. T1 x T2 tensor. durs (torch.tensor): List of durations corresponding to each text token. T2 tensor. Should sum to T1. start_token (int): Index of the starting token for the word of interest. num_tokens (int): Length (in tokens) of the word of interest. Output: mean_dist_for_word (float): Mean embedding distance between the word indicated and its predicted audio frames. """ # Need to calculate which audio frame we start on by summing all durations up to the start token's duration start_frame = torch.sum(durs[:start_token]).data total_frames = 0 dist_sum = 0 # Loop through each text token for token_ind in range(start_token, start_token + num_tokens): # Loop through each frame for the given text token for frame_ind in range(start_frame, start_frame + durs[token_ind]): # Recall that the L2 distance matrix is shape [spec_len, text_len] dist_sum += l2_dists[frame_ind, token_ind] # Update total frames so far & the starting frame for the next token total_frames += durs[token_ind] start_frame += durs[token_ind] return dist_sum / total_frames def forward(self, queries, keys, mask=None, attn_prior=None, conditioning=None): """Forward pass of the aligner encoder. Args: queries (torch.tensor): B x C x T1 tensor (probably going to be mel data). keys (torch.tensor): B x C2 x T2 tensor (text data). mask (torch.tensor): B x T2 x 1 tensor, binary mask for variable length entries (True = mask element, False = leave unchanged). attn_prior (torch.tensor): prior for attention matrix. conditioning (torch.tensor): B x T2 x 1 conditioning embedding Output: attn (torch.tensor): B x 1 x T1 x T2 attention mask. Final dim T2 should sum to 1. attn_logprob (torch.tensor): B x 1 x T1 x T2 log-prob attention mask. """ if conditioning is not None: keys = keys + conditioning.transpose(1, 2) keys_enc = self.key_proj(keys) # B x n_attn_dims x T2 queries_enc = self.query_proj(queries) # B x n_attn_dims x T1 # Simplistic Gaussian Isotopic Attention attn = (queries_enc[:, :, :, None] - keys_enc[:, :, None]) ** 2 # B x n_attn_dims x T1 x T2 attn = -self.temperature * attn.sum(1, keepdim=True) if attn_prior is not None: attn = self.log_softmax(attn) + torch.log(attn_prior[:, None] + 1e-8) attn_logprob = attn.clone() if mask is not None: attn.data.masked_fill_(mask.permute(0, 2, 1).unsqueeze(2), -float("inf")) attn = self.softmax(attn) # softmax along T2 return attn, attn_logprob class Encoder(nn.Module): """Encoder module: - Three 1-d convolution banks - Bidirectional LSTM """ def __init__(self, config): super(Encoder, self).__init__() convolutions = [] for _ in range(config["encoder"]["encoder_n_convolutions"]): conv_layer = nn.Sequential( ConvNorm(config["encoder"]["encoder_embedding_dim"], config["encoder"]["encoder_embedding_dim"], kernel_size=config["encoder"]["encoder_kernel_size"], stride=1, padding=int((config["encoder"]["encoder_kernel_size"] - 1) / 2), dilation=1, w_init_gain='relu'), nn.BatchNorm1d(config["encoder"]["encoder_embedding_dim"])) convolutions.append(conv_layer) self.convolutions = nn.ModuleList(convolutions) self.lstm = nn.LSTM(config["encoder"]["encoder_embedding_dim"], int(config["encoder"]["encoder_embedding_dim"] / 2), 1, batch_first=True, bidirectional=True) def forward(self, x, input_lengths): for conv in self.convolutions: x = F.dropout(F.relu(conv(x)), 0.5, self.training) x = x.transpose(1, 2) # pytorch tensor are not reversible, hence the conversion input_lengths = input_lengths.cpu().numpy() x = nn.utils.rnn.pack_padded_sequence( x, input_lengths, batch_first=True) self.lstm.flatten_parameters() outputs, _ = self.lstm(x) outputs, _ = nn.utils.rnn.pad_packed_sequence( outputs, batch_first=True) return outputs def inference(self, x): for conv in self.convolutions: x = F.dropout(F.relu(conv(x)), 0.5, self.training) x = x.transpose(1, 2) self.lstm.flatten_parameters() outputs, _ = self.lstm(x) return outputs class Decoder(nn.Module): def __init__(self, preprocess_config, model_config): super(Decoder, self).__init__() self.n_mel_channels = preprocess_config["preprocessing"]["mel"]["n_mel_channels"] self.n_frames_per_step = model_config["decoder"]["n_frames_per_step"] self.encoder_embedding_dim = model_config["encoder"]["encoder_embedding_dim"] self.attention_rnn_dim = model_config["attention"]["attention_rnn_dim"] self.decoder_rnn_dim = model_config["decoder"]["decoder_rnn_dim"] self.prenet_dim = model_config["decoder"]["prenet_dim"] self.max_decoder_steps = model_config["decoder"]["max_decoder_steps"] self.gate_threshold = model_config["decoder"]["gate_threshold"] self.p_attention_dropout = model_config["decoder"]["p_attention_dropout"] self.p_decoder_dropout = model_config["decoder"]["p_decoder_dropout"] attention_dim = model_config["attention"]["attention_dim"] attention_location_n_filters = model_config["location_layer"]["attention_location_n_filters"] attention_location_kernel_size = model_config["location_layer"]["attention_location_kernel_size"] self.prenet = Prenet( self.n_mel_channels * self.n_frames_per_step, [self.prenet_dim, self.prenet_dim]) self.attention_rnn = nn.LSTMCell( self.prenet_dim + self.encoder_embedding_dim, self.attention_rnn_dim) self.attention_layer = Attention( self.attention_rnn_dim, self.encoder_embedding_dim, attention_dim, attention_location_n_filters, attention_location_kernel_size) self.decoder_rnn = nn.LSTMCell( self.attention_rnn_dim + self.encoder_embedding_dim, self.decoder_rnn_dim, 1) self.linear_projection = LinearNorm( self.decoder_rnn_dim + self.encoder_embedding_dim, self.n_mel_channels * self.n_frames_per_step) self.gate_layer = LinearNorm( self.decoder_rnn_dim + self.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) # Grouping multiple frames if necessary: (B, n_mel_channels, T_out) -> (B, T_out/r, n_mel_channels*r) decoder_inputs = decoder_inputs.contiguous().view( decoder_inputs.size(0), int(decoder_inputs.size(1)/self.n_frames_per_step), -1) # (B, T_out/r, n_mel_channels*r) -> (T_out/r, B, n_mel_channels*r) 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/r, B) -> (B, T_out/r) alignments = torch.stack(alignments).transpose(0, 1) # (T_out/r, B) -> (B, T_out/r) gate_outputs = torch.stack(gate_outputs).transpose(0, 1) # tile gate_outputs to make frames per step. B = gate_outputs.size(0) gate_outputs = gate_outputs.contiguous().view(-1, 1).repeat(1,self.n_frames_per_step).view(B, -1) # (T_out/r, B, n_mel_channels*r) -> (B, T_out/r, n_mel_channels*r) mel_outputs = torch.stack(mel_outputs).transpose(0, 1).contiguous() # decouple frames per step: (B, T_out/r, n_mel_channels*r) -> (B, T_out, n_mel_channels) 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: """ cell_input = torch.cat((decoder_input, self.attention_context), -1) 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) attention_weights_cat = torch.cat( (self.attention_weights.unsqueeze(1), self.attention_weights_cum.unsqueeze(1)), dim=1) self.attention_context, self.attention_weights = self.attention_layer( self.attention_hidden, self.memory, self.processed_memory, attention_weights_cat, self.mask) 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) 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 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) # (1, B, n_mel_channels) decoder_inputs = self.parse_decoder_inputs(decoder_inputs) # (T_out/r, B, n_mel_channels*r) decoder_inputs = torch.cat((decoder_input, decoder_inputs), dim=0) # (1+(T_out/r), B, n_mel_channels*r) decoder_inputs = self.prenet(decoder_inputs) # (1+(T_out/r), B, prenet_dim) 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) mel_outputs += [mel_output.squeeze(1)] gate_outputs += [gate_output.squeeze(1)] alignments += [attention_weights] 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 // self.n_frames_per_step: 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 class LocationLayer(nn.Module): def __init__(self, attention_n_filters, attention_kernel_size, attention_dim): super(LocationLayer, self).__init__() padding = int((attention_kernel_size - 1) / 2) self.location_conv = ConvNorm(2, attention_n_filters, kernel_size=attention_kernel_size, padding=padding, bias=False, stride=1, dilation=1) self.location_dense = LinearNorm(attention_n_filters, attention_dim, bias=False, w_init_gain='tanh') def forward(self, attention_weights_cat): processed_attention = self.location_conv(attention_weights_cat) processed_attention = processed_attention.transpose(1, 2) processed_attention = self.location_dense(processed_attention) return processed_attention class Attention(nn.Module): def __init__(self, attention_rnn_dim, embedding_dim, attention_dim, attention_location_n_filters, attention_location_kernel_size): super(Attention, self).__init__() self.query_layer = LinearNorm(attention_rnn_dim, attention_dim, bias=False, w_init_gain='tanh') self.memory_layer = LinearNorm(embedding_dim, attention_dim, bias=False, w_init_gain='tanh') self.v = LinearNorm(attention_dim, 1, bias=False) self.location_layer = LocationLayer(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) 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) 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] self.layers = nn.ModuleList( [LinearNorm(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 Postnet(nn.Module): """Postnet - Five 1-d convolution with 512 channels and kernel size 5 """ def __init__(self, preprocess_config, model_config): super(Postnet, self).__init__() n_mel_channels = preprocess_config["preprocessing"]["mel"]["n_mel_channels"] postnet_embedding_dim = model_config["postnet"]["postnet_embedding_dim"] postnet_kernel_size = model_config["postnet"]["postnet_kernel_size"] postnet_n_convolutions = model_config["postnet"]["postnet_n_convolutions"] 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): 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) return x