mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-12 23:48:57 +00:00
Brandon Thomas
843 lines
33 KiB
Python
843 lines
33 KiB
Python
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
|