From 14356d32501beeb0e1a0a6713b88f3f936cf3292 Mon Sep 17 00:00:00 2001 From: erogol Date: Sat, 15 Aug 2020 00:17:49 +0200 Subject: [PATCH] glow-tts with relative pos encoding --- TTS/tts/utils/generic_utils.py | 2 + mozilla_voice_tts/bin/train_glow_tts.py | 26 +- .../tts/layers/glow_tts/attention.py | 202 ----------- .../tts/layers/glow_tts/commons.py | 63 ---- .../tts/layers/glow_tts/duration_predictor.py | 10 +- .../tts/layers/glow_tts/encoder.py | 65 ++-- .../tts/layers/glow_tts/glow_v2.py | 125 ------- .../tts/layers/glow_tts/transformer.py | 318 +++++++++++++++--- mozilla_voice_tts/tts/models/glow_tts.py | 6 + 9 files changed, 319 insertions(+), 498 deletions(-) delete mode 100644 mozilla_voice_tts/tts/layers/glow_tts/attention.py delete mode 100644 mozilla_voice_tts/tts/layers/glow_tts/commons.py delete mode 100644 mozilla_voice_tts/tts/layers/glow_tts/glow_v2.py diff --git a/TTS/tts/utils/generic_utils.py b/TTS/tts/utils/generic_utils.py index 2600e369..6d7ceb75 100644 --- a/TTS/tts/utils/generic_utils.py +++ b/TTS/tts/utils/generic_utils.py @@ -121,6 +121,8 @@ def setup_model(num_chars, num_speakers, c, speaker_embedding_dim=None): num_splits=4, num_sqz=2, sigmoid_scale=False, + rel_attn_window_size=4, + input_length=None, mean_only=True, hidden_channels_enc=192, hidden_channels_dec=192, diff --git a/mozilla_voice_tts/bin/train_glow_tts.py b/mozilla_voice_tts/bin/train_glow_tts.py index 72b8a7f1..186f2c18 100644 --- a/mozilla_voice_tts/bin/train_glow_tts.py +++ b/mozilla_voice_tts/bin/train_glow_tts.py @@ -18,25 +18,29 @@ from mozilla_voice_tts.tts.datasets.TTSDataset import MyDataset from mozilla_voice_tts.tts.layers.losses import GlowTTSLoss from mozilla_voice_tts.utils.console_logger import ConsoleLogger from mozilla_voice_tts.tts.utils.distribute import (DistributedSampler, - init_distributed, reduce_tensor) + init_distributed, + reduce_tensor) from mozilla_voice_tts.tts.utils.generic_utils import check_config, setup_model from mozilla_voice_tts.tts.utils.io import save_best_model, save_checkpoint from mozilla_voice_tts.tts.utils.measures import alignment_diagonal_score -from mozilla_voice_tts.tts.utils.speakers import (get_speakers, load_speaker_mapping, - save_speaker_mapping) +from mozilla_voice_tts.tts.utils.speakers import (get_speakers, + load_speaker_mapping, + save_speaker_mapping) from mozilla_voice_tts.tts.utils.synthesis import synthesis from mozilla_voice_tts.tts.utils.text.symbols import make_symbols, phonemes, symbols from mozilla_voice_tts.tts.utils.visual import plot_alignment, plot_spectrogram from mozilla_voice_tts.utils.audio import AudioProcessor -from mozilla_voice_tts.utils.generic_utils import (KeepAverage, count_parameters, - create_experiment_folder, get_git_branch, - remove_experiment_folder, set_init_dict) +from mozilla_voice_tts.utils.generic_utils import ( + KeepAverage, count_parameters, create_experiment_folder, get_git_branch, + remove_experiment_folder, set_init_dict) from mozilla_voice_tts.utils.io import copy_config_file, load_config from mozilla_voice_tts.utils.radam import RAdam from mozilla_voice_tts.utils.tensorboard_logger import TensorboardLogger -from mozilla_voice_tts.utils.training import (NoamLR, adam_weight_decay, check_update, - gradual_training_scheduler, set_weight_decay, - setup_torch_training_env) +from mozilla_voice_tts.utils.training import (NoamLR, adam_weight_decay, + check_update, + gradual_training_scheduler, + set_weight_decay, + setup_torch_training_env) use_cuda, num_gpus = setup_torch_training_env(True, False) @@ -118,7 +122,7 @@ def data_depended_init(model, ap): if getattr(f, "set_ddi", False): f.set_ddi(True) else: - for f in model.decoder.flows: + for f in model.decoder.flows: if getattr(f, "set_ddi", False): f.set_ddi(True) @@ -142,7 +146,7 @@ def data_depended_init(model, ap): if getattr(f, "set_ddi", False): f.set_ddi(False) else: - for f in model.decoder.flows: + for f in model.decoder.flows: if getattr(f, "set_ddi", False): f.set_ddi(False) return model diff --git a/mozilla_voice_tts/tts/layers/glow_tts/attention.py b/mozilla_voice_tts/tts/layers/glow_tts/attention.py deleted file mode 100644 index f7e14ddc..00000000 --- a/mozilla_voice_tts/tts/layers/glow_tts/attention.py +++ /dev/null @@ -1,202 +0,0 @@ -import copy -import math -import numpy as np -import torch -from torch import nn -from torch.nn import functional as F - - -def convert_pad_shape(pad_shape): - l = pad_shape[::-1] - pad_shape = [item for sublist in l for item in sublist] - return pad_shape - - -class MultiHeadAttention(nn.Module): - def __init__(self, - channels, - out_channels, - n_heads, - window_size=None, - heads_share=True, - dropout_p=0., - input_length=None, - proximal_bias=False, - proximal_init=False): - super().__init__() - assert channels % n_heads == 0 - - self.channels = channels - self.out_channels = out_channels - self.n_heads = n_heads - self.window_size = window_size - self.heads_share = heads_share - self.input_length = input_length - self.proximal_bias = proximal_bias - self.dropout_p = dropout_p - self.attn = None - - self.k_channels = channels // n_heads - self.conv_q = nn.Conv1d(channels, channels, 1) - self.conv_k = nn.Conv1d(channels, channels, 1) - self.conv_v = nn.Conv1d(channels, channels, 1) - if window_size is not None: - n_heads_rel = 1 if heads_share else n_heads - rel_stddev = self.k_channels**-0.5 - self.emb_rel_k = nn.Parameter( - torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) - * rel_stddev) - self.emb_rel_v = nn.Parameter( - torch.randn(n_heads_rel, window_size * 2 + 1, self.k_channels) - * rel_stddev) - self.conv_o = nn.Conv1d(channels, out_channels, 1) - self.drop = nn.Dropout(dropout_p) - - nn.init.xavier_uniform_(self.conv_q.weight) - nn.init.xavier_uniform_(self.conv_k.weight) - if proximal_init: - self.conv_k.weight.data.copy_(self.conv_q.weight.data) - self.conv_k.bias.data.copy_(self.conv_q.bias.data) - nn.init.xavier_uniform_(self.conv_v.weight) - - def forward(self, x, c, attn_mask=None): - q = self.conv_q(x) - k = self.conv_k(c) - v = self.conv_v(c) - - x, self.attn = self.attention(q, k, v, mask=attn_mask) - - x = self.conv_o(x) - return x - - def attention(self, query, key, value, mask=None): - # reshape [b, d, t] -> [b, n_h, t, d_k] - b, d, t_s, t_t = (*key.size(), query.size(2)) - query = query.view(b, self.n_heads, self.k_channels, - t_t).transpose(2, 3) - key = key.view(b, self.n_heads, self.k_channels, t_s).transpose(2, 3) - value = value.view(b, self.n_heads, self.k_channels, - t_s).transpose(2, 3) - - scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt( - self.k_channels) - if self.window_size is not None: - assert t_s == t_t, "Relative attention is only available for self-attention." - key_relative_embeddings = self._get_relative_embeddings( - self.emb_rel_k, t_s) - rel_logits = self._matmul_with_relative_keys( - query, key_relative_embeddings) - rel_logits = self._relative_position_to_absolute_position( - rel_logits) - scores_local = rel_logits / math.sqrt(self.k_channels) - scores = scores + scores_local - if self.proximal_bias: - assert t_s == t_t, "Proximal bias is only available for self-attention." - scores = scores + self._attention_bias_proximal(t_s).to( - device=scores.device, dtype=scores.dtype) - if mask is not None: - scores = scores.masked_fill(mask == 0, -1e4) - if self.input_length is not None: - block_mask = torch.ones_like(scores).triu( - -self.input_length).tril(self.input_length) - scores = scores * block_mask + -1e4 * (1 - block_mask) - p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s] - p_attn = self.drop(p_attn) - output = torch.matmul(p_attn, value) - if self.window_size is not None: - relative_weights = self._absolute_position_to_relative_position( - p_attn) - value_relative_embeddings = self._get_relative_embeddings( - self.emb_rel_v, t_s) - output = output + self._matmul_with_relative_values( - relative_weights, value_relative_embeddings) - output = output.transpose(2, 3).contiguous().view( - b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t] - return output, p_attn - - def _matmul_with_relative_values(self, x, y): - """ - x: [b, h, l, m] - y: [h or 1, m, d] - ret: [b, h, l, d] - """ - ret = torch.matmul(x, y.unsqueeze(0)) - return ret - - def _matmul_with_relative_keys(self, x, y): - """ - x: [b, h, l, d] - y: [h or 1, m, d] - ret: [b, h, l, m] - """ - ret = torch.matmul(x, y.unsqueeze(0).transpose(-2, -1)) - return ret - - def _get_relative_embeddings(self, relative_embeddings, length): - max_relative_position = 2 * self.window_size + 1 - # Pad first before slice to avoid using cond ops. - pad_length = max(length - (self.window_size + 1), 0) - slice_start_position = max((self.window_size + 1) - length, 0) - slice_end_position = slice_start_position + 2 * length - 1 - if pad_length > 0: - padded_relative_embeddings = F.pad( - relative_embeddings, - convert_pad_shape([[0, 0], [pad_length, pad_length], - [0, 0]])) - else: - padded_relative_embeddings = relative_embeddings - used_relative_embeddings = padded_relative_embeddings[:, - slice_start_position: - slice_end_position] - return used_relative_embeddings - - def _relative_position_to_absolute_position(self, x): - """ - x: [b, h, l, 2*l-1] - ret: [b, h, l, l] - """ - batch, heads, length, _ = x.size() - # Concat columns of pad to shift from relative to absolute indexing. - x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, - 1]])) - - # Concat extra elements so to add up to shape (len+1, 2*len-1). - x_flat = x.view([batch, heads, length * 2 * length]) - x_flat = F.pad( - x_flat, convert_pad_shape([[0, 0], [0, 0], [0, - length - 1]])) - - # Reshape and slice out the padded elements. - x_final = x_flat.view([batch, heads, length + 1, - 2 * length - 1])[:, :, :length, length - 1:] - return x_final - - def _absolute_position_to_relative_position(self, x): - """ - x: [b, h, l, l] - ret: [b, h, l, 2*l-1] - """ - batch, heads, length, _ = x.size() - # padd along column - x = F.pad( - x, - convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, - length - 1]])) - x_flat = x.view([batch, heads, length**2 + length * (length - 1)]) - # add 0's in the beginning that will skew the elements after reshape - x_flat = F.pad( - x_flat, convert_pad_shape([[0, 0], [0, 0], [length, 0]])) - x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:] - return x_final - - def _attention_bias_proximal(self, length): - """Bias for self-attention to encourage attention to close positions. - Args: - length: an integer scalar. - Returns: - a Tensor with shape [1, 1, length, length] - """ - r = torch.arange(length, dtype=torch.float32) - diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1) - return torch.unsqueeze( - torch.unsqueeze(-torch.log1p(torch.abs(diff)), 0), 0) diff --git a/mozilla_voice_tts/tts/layers/glow_tts/commons.py b/mozilla_voice_tts/tts/layers/glow_tts/commons.py deleted file mode 100644 index 5724a909..00000000 --- a/mozilla_voice_tts/tts/layers/glow_tts/commons.py +++ /dev/null @@ -1,63 +0,0 @@ -import math -import numpy as np -import torch -from torch import nn -from torch.nn import functional as F - - -def convert_pad_shape(pad_shape): - l = pad_shape[::-1] - pad_shape = [item for sublist in l for item in sublist] - return pad_shape - - -def shift_1d(x): - x = F.pad(x, convert_pad_shape([[0, 0], [0, 0], [1, 0]]))[:, :, :-1] - return x - - -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) - - -def maximum_path(value, mask, max_neg_val=-np.inf): - """ Numpy-friendly version. It's about 4 times faster than torch version. - value: [b, t_x, t_y] - mask: [b, t_x, t_y] - """ - value = value * mask - - device = value.device - dtype = value.dtype - value = value.cpu().detach().numpy() - mask = mask.cpu().detach().numpy().astype(np.bool) - - b, t_x, t_y = value.shape - direction = np.zeros(value.shape, dtype=np.int64) - v = np.zeros((b, t_x), dtype=np.float32) - x_range = np.arange(t_x, dtype=np.float32).reshape(1, -1) - for j in range(t_y): - v0 = np.pad(v, [[0, 0], [1, 0]], - mode="constant", - constant_values=max_neg_val)[:, :-1] - v1 = v - max_mask = (v1 >= v0) - v_max = np.where(max_mask, v1, v0) - direction[:, :, j] = max_mask - - index_mask = (x_range <= j) - v = np.where(index_mask, v_max + value[:, :, j], max_neg_val) - direction = np.where(mask, direction, 1) - - path = np.zeros(value.shape, dtype=np.float32) - index = mask[:, :, 0].sum(1).astype(np.int64) - 1 - index_range = np.arange(b) - for j in reversed(range(t_y)): - path[index_range, index, j] = 1 - index = index + direction[index_range, index, j] - 1 - path = path * mask.astype(np.float32) - path = torch.from_numpy(path).to(device=device, dtype=dtype) - return path diff --git a/mozilla_voice_tts/tts/layers/glow_tts/duration_predictor.py b/mozilla_voice_tts/tts/layers/glow_tts/duration_predictor.py index 4b17177e..5a321662 100644 --- a/mozilla_voice_tts/tts/layers/glow_tts/duration_predictor.py +++ b/mozilla_voice_tts/tts/layers/glow_tts/duration_predictor.py @@ -1,21 +1,16 @@ -import copy -import math -import numpy as np -import scipy import torch from torch import nn -from torch.nn import functional as F class DurationPredictor(nn.Module): def __init__(self, in_channels, filter_channels, kernel_size, dropout_p): super().__init__() - + # class arguments self.in_channels = in_channels self.filter_channels = filter_channels self.kernel_size = kernel_size self.dropout_p = dropout_p - + # layers self.drop = nn.Dropout(dropout_p) self.conv_1 = nn.Conv1d(in_channels, filter_channels, @@ -27,6 +22,7 @@ class DurationPredictor(nn.Module): kernel_size, padding=kernel_size // 2) self.norm_2 = nn.GroupNorm(1, filter_channels) + # output layer self.proj = nn.Conv1d(filter_channels, 1, 1) def forward(self, x, x_mask): diff --git a/mozilla_voice_tts/tts/layers/glow_tts/encoder.py b/mozilla_voice_tts/tts/layers/glow_tts/encoder.py index d8bbc07c..b0e40ca4 100644 --- a/mozilla_voice_tts/tts/layers/glow_tts/encoder.py +++ b/mozilla_voice_tts/tts/layers/glow_tts/encoder.py @@ -1,33 +1,13 @@ import math import torch from torch import nn -from torch.nn import functional as F +from mozilla_voice_tts.tts.layers.glow_tts.transformer import Transformer from mozilla_voice_tts.tts.utils.generic_utils import sequence_mask from mozilla_voice_tts.tts.layers.glow_tts.glow import ConvLayerNorm from mozilla_voice_tts.tts.layers.glow_tts.duration_predictor import DurationPredictor -class PositionalEncoding(nn.Module): - def __init__(self, d_model, max_len=5000): - super(PositionalEncoding, self).__init__() - self.register_buffer('pe', self._get_pe_matrix(d_model, max_len)) - - def forward(self, x): - return x + self.pe[:x.size(0)].unsqueeze(1) - - def _get_pe_matrix(self, d_model, max_len): - pe = torch.zeros(max_len, d_model) - position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) - div_term = torch.pow(10000, - torch.arange(0, d_model, 2).float() / d_model) - - pe[:, 0::2] = torch.sin(position / div_term) - pe[:, 1::2] = torch.cos(position / div_term) - - return pe - - class Encoder(nn.Module): """Glow-TTS encoder module. We use Pytorch TransformerEncoder instead of the one with relative position embedding. We use positional encoding @@ -60,12 +40,13 @@ class Encoder(nn.Module): num_layers, kernel_size, dropout_p, + rel_attn_window_size=None, + input_length=None, mean_only=False, use_prenet=False, c_in_channels=0): - super().__init__() - + # class arguments self.num_chars = num_chars self.out_channels = out_channels self.hidden_channels = hidden_channels @@ -78,11 +59,10 @@ class Encoder(nn.Module): self.mean_only = mean_only self.use_prenet = use_prenet self.c_in_channels = c_in_channels - + # embedding layer self.emb = nn.Embedding(num_chars, hidden_channels) nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5) - self.register_buffer('pe', PositionalEncoding(hidden_channels).pe) - + # optional convolutional prenet if use_prenet: self.pre = ConvLayerNorm(hidden_channels, hidden_channels, @@ -90,49 +70,50 @@ class Encoder(nn.Module): kernel_size=5, num_layers=3, dropout_p=0.5) - - encoder = nn.TransformerEncoderLayer(hidden_channels, - num_heads, - dim_feedforward=filter_channels, - dropout=dropout_p) - self.encoder = nn.TransformerEncoder(encoder, num_layers) - + # text encoder + self.encoder = Transformer(hidden_channels, + filter_channels, + num_heads, + num_layers, + kernel_size=kernel_size, + dropout_p=dropout_p, + rel_attn_window_size=rel_attn_window_size, + input_length=input_length) + # final projection layers self.proj_m = nn.Conv1d(hidden_channels, out_channels, 1) if not mean_only: self.proj_s = nn.Conv1d(hidden_channels, out_channels, 1) + # duration predictor self.duration_predictor = DurationPredictor( hidden_channels + c_in_channels, filter_channels_dp, kernel_size, dropout_p) def forward(self, x, x_lengths, g=None): - # pass embedding layer + # embedding layer # [B ,T, D] x = self.emb(x) * math.sqrt(self.hidden_channels) - x = x + self.pe[:x.shape[1]].unsqueeze(0) # [B, D, T] x = torch.transpose(x, 1, -1) # compute input sequence mask x_mask = torch.unsqueeze(sequence_mask(x_lengths, x.size(2)), 1).to(x.dtype) - # pass pre-conv layers + # pre-conv layers if self.use_prenet: x = self.pre(x, x_mask) - # pass encoder - x = self.encoder(x.permute(2, 0, 1), - src_key_padding_mask=~(x_mask.squeeze(1).to(bool))) - x = x.permute(1, 2, 0) + # encoder + x = self.encoder(x, x_mask) # set duration predictor input if g is not None: g_exp = g.expand(-1, -1, x.size(-1)) x_dp = torch.cat([torch.detach(x), g_exp], 1) else: x_dp = torch.detach(x) - # pass final projection layer + # final projection layer x_m = self.proj_m(x) * x_mask if not self.mean_only: x_logs = self.proj_s(x) * x_mask else: x_logs = torch.zeros_like(x_m) - # pass duration predictor + # duration predictor logw = self.duration_predictor(x_dp, x_mask) return x_m, x_logs, logw, x_mask diff --git a/mozilla_voice_tts/tts/layers/glow_tts/glow_v2.py b/mozilla_voice_tts/tts/layers/glow_tts/glow_v2.py deleted file mode 100644 index 28bf960f..00000000 --- a/mozilla_voice_tts/tts/layers/glow_tts/glow_v2.py +++ /dev/null @@ -1,125 +0,0 @@ -import torch -from torch import nn -from torch.nn import functional as F - -from mozilla_voice_tts.tts.utils.generic_utils import sequence_mask -from mozilla_voice_tts.tts.layers.glow_tts.glow import InvConvNear, CouplingBlock, ActNorm, BatchNorm - - -def squeeze(x, x_mask=None, num_sqz=2): - b, c, t = x.size() - - t = (t // num_sqz) * num_sqz - x = x[:, :, :t] - x_sqz = x.view(b, c, t // num_sqz, num_sqz) - x_sqz = x_sqz.permute(0, 3, 1, - 2).contiguous().view(b, c * num_sqz, t // num_sqz) - - if x_mask is not None: - x_mask = x_mask[:, :, num_sqz - 1::num_sqz] - else: - x_mask = torch.ones(b, 1, t // num_sqz).to(device=x.device, - dtype=x.dtype) - return x_sqz * x_mask, x_mask - - -def unsqueeze(x, x_mask=None, num_sqz=2): - b, c, t = x.size() - - x_unsqz = x.view(b, num_sqz, c // num_sqz, t) - x_unsqz = x_unsqz.permute(0, 2, 3, - 1).contiguous().view(b, c // num_sqz, - t * num_sqz) - - if x_mask is not None: - x_mask = x_mask.unsqueeze(-1).repeat(1, 1, 1, - num_sqz).view(b, 1, t * num_sqz) - else: - x_mask = torch.ones(b, 1, t * num_sqz).to(device=x.device, - dtype=x.dtype) - return x_unsqz * x_mask, x_mask - - -class Decoder(nn.Module): - """Stack of Glow Modules""" - def __init__(self, - in_channels, - hidden_channels, - kernel_size, - dilation_rate, - num_blocks, - num_coupling_layers, - dropout_p=0., - num_splits=4, - num_sqz=2, - sigmoid_scale=False, - c_in_channels=0): - super().__init__() - - self.in_channels = in_channels - self.hidden_channels = hidden_channels - self.kernel_size = kernel_size - self.dilation_rate = dilation_rate - self.num_blocks = num_blocks - self.num_coupling_layers = num_coupling_layers - self.dropout_p = dropout_p - self.num_splits = num_splits - self.num_sqz = num_sqz - self.sigmoid_scale = sigmoid_scale - self.c_in_channels = c_in_channels - - self.norm_layers = nn.ModuleList() - self.invconv_layers = nn.ModuleList() - self.coupling_layers = nn.ModuleList() - - for _ in range(num_blocks): - self.norm_layers.append(ActNorm(in_channels=in_channels * num_sqz)) - self.invconv_layers.append( - InvConvNear(channels=in_channels * num_sqz, - num_splits=num_splits)) - self.coupling_layers.append( - CouplingBlock(in_channels * num_sqz, - hidden_channels, - kernel_size=kernel_size, - dilation_rate=dilation_rate, - num_layers=num_coupling_layers, - c_in_channels=c_in_channels, - dropout_p=dropout_p, - sigmoid_scale=sigmoid_scale)) - - def forward(self, x, x_mask, g=None, reverse=False): - if not reverse: - # norm_layers = self.norm_layers - invconv_layers = self.invconv_layers - coupling_layers = self.coupling_layers - logdet_tot = 0 - else: - # norm_layers = reversed(self.norm_layers) - invconv_layers = self.invconv_layers[::-1] - coupling_layers = self.coupling_layers[::-1] - logdet_tot = None - - if self.num_sqz > 1: - x, x_mask = squeeze(x, x_mask, self.num_sqz) - for idx in range(len(self.invconv_layers)): - if not reverse: - # x, log_det = norm_layers[idx](x, reverse) - # logdet_tot += log_det - - x, log_det = invconv_layers[idx](x, x_mask, reverse) - logdet_tot += log_det - - x, log_det = coupling_layers[idx](x, x_mask, g=g, reverse=reverse) - logdet_tot += log_det - else: - x, log_det = coupling_layers[idx](x, x_mask, g=g, reverse=reverse) - x, log_det = invconv_layers[idx](x, x_mask, reverse) - # x, logdet = norm_layers[idx](x, reverse) - - if self.num_sqz > 1: - x, x_mask = unsqueeze(x, x_mask, self.num_sqz) - return x, logdet_tot - - def store_inverse(self): - for f in self.flows: - f.store_inverse() \ No newline at end of file diff --git a/mozilla_voice_tts/tts/layers/glow_tts/transformer.py b/mozilla_voice_tts/tts/layers/glow_tts/transformer.py index 0081130e..366e48a8 100644 --- a/mozilla_voice_tts/tts/layers/glow_tts/transformer.py +++ b/mozilla_voice_tts/tts/layers/glow_tts/transformer.py @@ -1,67 +1,232 @@ +import copy +import math +import numpy as np import torch from torch import nn from torch.nn import functional as F -from mozilla_voice_tts.tts.layers.glow_tts.attention import MultiHeadAttention +from mozilla_voice_tts.tts.layers.glow_tts.glow import LayerNorm -class Transformer(nn.Module): +class RelativePositionMultiHeadAttention(nn.Module): + """Implementation of https://arxiv.org/pdf/1803.02155.pdf + Visualization of the algorithm + https://raw.githubusercontent.com/Separius/CudaRelativeAttention/master/algorithm.png + """ def __init__(self, - hidden_channels, - filter_channels, + channels, + out_channels, num_heads, - num_layers, - kernel_size=1, - dropout_p=0., rel_attn_window_size=None, + heads_share=True, + dropout_p=0., input_length=None, - **kwargs): + proximal_bias=False, + proximal_init=False): super().__init__() - self.hidden_channels = hidden_channels - self.filter_channels = filter_channels + assert channels % num_heads == 0, " [!] channels should be divisible by num_heads." + # class attributes + self.channels = channels + self.out_channels = out_channels self.num_heads = num_heads - self.num_layers = num_layers - self.kernel_size = kernel_size - self.dropout_p = dropout_p self.rel_attn_window_size = rel_attn_window_size + self.heads_share = heads_share self.input_length = input_length + self.proximal_bias = proximal_bias + self.dropout_p = dropout_p + self.attn = None + # query, key, value layers + self.k_channels = channels // num_heads + self.conv_q = nn.Conv1d(channels, channels, 1) + self.conv_k = nn.Conv1d(channels, channels, 1) + self.conv_v = nn.Conv1d(channels, channels, 1) + # output layers + self.conv_o = nn.Conv1d(channels, out_channels, 1) + self.dropout = nn.Dropout(dropout_p) + # relative positional encoding layers + if rel_attn_window_size is not None: + n_heads_rel = 1 if heads_share else num_heads + rel_stddev = self.k_channels**-0.5 + emb_rel_k = nn.Parameter( + torch.randn(n_heads_rel, rel_attn_window_size * 2 + 1, + self.k_channels) * rel_stddev) + emb_rel_v = nn.Parameter( + torch.randn(n_heads_rel, rel_attn_window_size * 2 + 1, + self.k_channels) * rel_stddev) + self.register_parameter('emb_rel_k', emb_rel_k) + self.register_parameter('emb_rel_v', emb_rel_v) - self.drop = nn.Dropout(dropout_p) - self.attn_layers = nn.ModuleList() - self.norm_layers_1 = nn.ModuleList() - self.ffn_layers = nn.ModuleList() - self.norm_layers_2 = nn.ModuleList() - for _ in range(self.num_layers): - self.attn_layers.append( - MultiHeadAttention(hidden_channels, - hidden_channels, - num_heads, - window_size=rel_attn_window_size, - dropout_p=dropout_p, - input_length=input_length)) - self.norm_layers_1.append(nn.GroupNorm(1, hidden_channels)) - self.ffn_layers.append( - FFN(hidden_channels, - hidden_channels, - filter_channels, - kernel_size, - dropout_p=dropout_p)) - self.norm_layers_2.append(nn.GroupNorm(1, hidden_channels)) + # init layers + nn.init.xavier_uniform_(self.conv_q.weight) + nn.init.xavier_uniform_(self.conv_k.weight) + # proximal bias + if proximal_init: + self.conv_k.weight.data.copy_(self.conv_q.weight.data) + self.conv_k.bias.data.copy_(self.conv_q.bias.data) + nn.init.xavier_uniform_(self.conv_v.weight) - def forward(self, x, x_mask): - attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1) - for i in range(self.num_layers): - x = x * x_mask - y = self.attn_layers[i](x, x, attn_mask) - y = self.drop(y) - x = self.norm_layers_1[i](x + y) - - y = self.ffn_layers[i](x, x_mask) - y = self.drop(y) - x = self.norm_layers_2[i](x + y) - x = x * x_mask + def forward(self, x, c, attn_mask=None): + q = self.conv_q(x) + k = self.conv_k(c) + v = self.conv_v(c) + x, self.attn = self.attention(q, k, v, mask=attn_mask) + x = self.conv_o(x) return x + def attention(self, query, key, value, mask=None): + # reshape [b, d, t] -> [b, n_h, t, d_k] + b, d, t_s, t_t = (*key.size(), query.size(2)) + query = query.view(b, self.num_heads, self.k_channels, + t_t).transpose(2, 3) + key = key.view(b, self.num_heads, self.k_channels, t_s).transpose(2, 3) + value = value.view(b, self.num_heads, self.k_channels, + t_s).transpose(2, 3) + # compute raw attention scores + scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt( + self.k_channels) + # relative positional encoding + if self.rel_attn_window_size is not None: + assert t_s == t_t, "Relative attention is only available for self-attention." + # get relative key embeddings + breakpoint() + key_relative_embeddings = self._get_relative_embeddings( + self.emb_rel_k, t_s) + rel_logits = self._matmul_with_relative_keys( + query, key_relative_embeddings) + rel_logits = self._relative_position_to_absolute_position( + rel_logits) + scores_local = rel_logits / math.sqrt(self.k_channels) + scores = scores + scores_local + # proximan bias + if self.proximal_bias: + assert t_s == t_t, "Proximal bias is only available for self-attention." + scores = scores + self._attn_proximity_bias(t_s).to( + device=scores.device, dtype=scores.dtype) + # attention score masking + if mask is not None: + # add small value to prevent oor error. + scores = scores.masked_fill(mask == 0, -1e4) + if self.input_length is not None: + block_mask = torch.ones_like(scores).triu( + -self.input_length).tril(self.input_length) + scores = scores * block_mask + -1e4 * (1 - block_mask) + # attention score normalization + p_attn = F.softmax(scores, dim=-1) # [b, n_h, t_t, t_s] + # apply dropout to attention weights + p_attn = self.dropout(p_attn) + # compute output + output = torch.matmul(p_attn, value) + # relative positional encoding for values + if self.rel_attn_window_size is not None: + relative_weights = self._absolute_position_to_relative_position( + p_attn) + value_relative_embeddings = self._get_relative_embeddings( + self.emb_rel_v, t_s) + output = output + self._matmul_with_relative_values( + relative_weights, value_relative_embeddings) + output = output.transpose(2, 3).contiguous().view( + b, d, t_t) # [b, n_h, t_t, d_k] -> [b, d, t_t] + return output, p_attn + + def _matmul_with_relative_values(self, p_attn, re): + """ + Args: + p_attn (Tensor): attention weights. + re (Tensor): relative value embedding vector. (a_(i,j)^V) + + Shapes: + p_attn: [B, H, T, V] + re: [H or 1, V, D] + logits: [B, H, T, D] + """ + logits = torch.matmul(p_attn, re.unsqueeze(0)) + return logits + + @staticmethod + def _matmul_with_relative_keys(query, re): + """ + Args: + query (Tensor): batch of query vectors. (x*W^Q) + re (Tensor): relative key embedding vector. (a_(i,j)^K) + + Shapes: + query: [B, H, T, D] + re: [H or 1, V, D] + logits: [B, H, T, V] + """ + # logits = torch.einsum('bhld, kmd -> bhlm', [query, re.to(query.dtype)]) + logits = torch.matmul(query, re.unsqueeze(0).transpose(-2, -1)) + return logits + + def _get_relative_embeddings(self, relative_embeddings, length): + """Convert embedding vestors to a tensor of embeddings + """ + # Pad first before slice to avoid using cond ops. + pad_length = max(length - (self.rel_attn_window_size + 1), 0) + slice_start_position = max((self.rel_attn_window_size + 1) - length, 0) + slice_end_position = slice_start_position + 2 * length - 1 + if pad_length > 0: + padded_relative_embeddings = F.pad( + relative_embeddings, [0, 0, pad_length, pad_length, 0, 0]) + else: + padded_relative_embeddings = relative_embeddings + used_relative_embeddings = padded_relative_embeddings[:, + slice_start_position: + slice_end_position] + return used_relative_embeddings + + @staticmethod + def _relative_position_to_absolute_position(x): + """Converts tensor from relative to absolute indexing for local attention. + Args: + x: [B, D, length, 2 * length - 1] + Returns: + A Tensor of shape [B, D, length, length] + """ + batch, heads, length, _ = x.size() + # Pad to shift from relative to absolute indexing. + x = F.pad(x, [0, 1, 0, 0, 0, 0, 0, 0]) + # Pad extra elements so to add up to shape (len+1, 2*len-1). + x_flat = x.view([batch, heads, length * 2 * length]) + x_flat = F.pad(x_flat, [0, length - 1, 0, 0, 0, 0]) + # Reshape and slice out the padded elements. + x_final = x_flat.view([batch, heads, length + 1, + 2 * length - 1])[:, :, :length, length - 1:] + return x_final + + @staticmethod + def _absolute_position_to_relative_position(x): + """ + x: [B, H, T, T] + ret: [B, H, T, 2*T-1] + """ + batch, heads, length, _ = x.size() + # padd along column + x = F.pad(x, [0, length - 1, 0, 0, 0, 0, 0, 0]) + x_flat = x.view([batch, heads, length**2 + length * (length - 1)]) + # add 0's in the beginning that will skew the elements after reshape + x_flat = F.pad(x_flat, [length, 0, 0, 0, 0, 0]) + x_final = x_flat.view([batch, heads, length, 2 * length])[:, :, :, 1:] + return x_final + + @staticmethod + def _attn_proximity_bias(length): + """Produce an attention mask that discourages distant + attention values. + Args: + length (int): an integer scalar. + Returns: + a Tensor with shape [1, 1, length, length] + """ + # L + r = torch.arange(length, dtype=torch.float32) + # L x L + diff = torch.unsqueeze(r, 0) - torch.unsqueeze(r, 1) + # scale mask values + diff = -torch.log1p(torch.abs(diff)) + # 1 x 1 x L x L + return diff.unsqueeze(0).unsqueeze(0) + class FFN(nn.Module): def __init__(self, @@ -87,7 +252,7 @@ class FFN(nn.Module): out_channels, kernel_size, padding=kernel_size // 2) - self.drop = nn.Dropout(dropout_p) + self.dropout = nn.Dropout(dropout_p) def forward(self, x, x_mask): x = self.conv_1(x * x_mask) @@ -95,6 +260,63 @@ class FFN(nn.Module): x = x * torch.sigmoid(1.702 * x) else: x = torch.relu(x) - x = self.drop(x) + x = self.dropout(x) x = self.conv_2(x * x_mask) return x * x_mask + + +class Transformer(nn.Module): + def __init__(self, + hidden_channels, + filter_channels, + num_heads, + num_layers, + kernel_size=1, + dropout_p=0., + rel_attn_window_size=None, + input_length=None): + super().__init__() + self.hidden_channels = hidden_channels + self.filter_channels = filter_channels + self.num_heads = num_heads + self.num_layers = num_layers + self.kernel_size = kernel_size + self.dropout_p = dropout_p + self.rel_attn_window_size = rel_attn_window_size + + self.dropout = nn.Dropout(dropout_p) + self.attn_layers = nn.ModuleList() + self.norm_layers_1 = nn.ModuleList() + self.ffn_layers = nn.ModuleList() + self.norm_layers_2 = nn.ModuleList() + for _ in range(self.num_layers): + self.attn_layers.append( + RelativePositionMultiHeadAttention( + hidden_channels, + hidden_channels, + num_heads, + rel_attn_window_size=rel_attn_window_size, + dropout_p=dropout_p, + input_length=input_length)) + self.norm_layers_1.append(LayerNorm(hidden_channels)) + self.ffn_layers.append( + FFN(hidden_channels, + hidden_channels, + filter_channels, + kernel_size, + dropout_p=dropout_p)) + self.norm_layers_2.append(LayerNorm(hidden_channels)) + + def forward(self, x, x_mask): + attn_mask = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1) + for i in range(self.num_layers): + x = x * x_mask + y = self.attn_layers[i](x, x, attn_mask) + y = self.dropout(y) + x = self.norm_layers_1[i](x + y) + + y = self.ffn_layers[i](x, x_mask) + y = self.dropout(y) + x = self.norm_layers_2[i](x + y) + x = x * x_mask + return x diff --git a/mozilla_voice_tts/tts/models/glow_tts.py b/mozilla_voice_tts/tts/models/glow_tts.py index cac6ec1c..9bb96ae4 100644 --- a/mozilla_voice_tts/tts/models/glow_tts.py +++ b/mozilla_voice_tts/tts/models/glow_tts.py @@ -31,6 +31,8 @@ class GlowTts(nn.Module): num_splits=4, num_sqz=1, sigmoid_scale=False, + rel_attn_window_size=None, + input_length=None, mean_only=False, hidden_channels_enc=None, hidden_channels_dec=None, @@ -56,6 +58,8 @@ class GlowTts(nn.Module): self.num_splits = num_splits self.num_sqz = num_sqz self.sigmoid_scale = sigmoid_scale + self.rel_attn_window_size = rel_attn_window_size + self.input_length = input_length self.mean_only = mean_only self.hidden_channels_enc = hidden_channels_enc self.hidden_channels_dec = hidden_channels_dec @@ -72,6 +76,8 @@ class GlowTts(nn.Module): num_layers_enc, kernel_size, dropout_p, + rel_attn_window_size=rel_attn_window_size, + input_length=input_length, mean_only=mean_only, use_prenet=True, c_in_channels=c_in_channels)