tacotron.py

""" 以下代码仅为Tacotron模型的一个参考实现
"""

import torch
import torch.nn as nn

# Tacotron编码器
class Encoder(nn.Module):
    def __init__(self, encoder_n_convolutions,
        encoder_embedding_dim, encoder_kernel_size):
        super(Encoder, self).__init__()

        convolutions = []
        for _ in range(encoder_n_convolutions):
            conv_layer = nn.Sequential(
                nn.Conv1d(
                    encoder_embedding_dim,
                    encoder_embedding_dim,
                    kernel_size=encoder_kernel_size, 
                    stride=1,
                    padding=encoder_kernel_size//2,
                         dilation=1),
                nn.BatchNorm1d(encoder_embedding_dim))
            convolutions.append(conv_layer)
        self.convolutions = nn.ModuleList(convolutions)

        self.lstm = nn.LSTM(encoder_embedding_dim,
                            encoder_embedding_dim // 2, 1,
                            batch_first=True, bidirectional=True)

    def forward(self, x, input_lengths):
        # 假设输入为N×C×T
        for conv in self.convolutions:
            x = F.dropout(F.relu(conv(x)), 0.5, self.training)

        x = x.transpose(1, 2)

        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

# Tacotron前处理/后处理代码
class Prenet(nn.Module):
    def __init__(self, in_dim, sizes):
        super(Prenet, self).__init__()
        in_sizes = [in_dim] + sizes[:-1]
        self.layers = nn.ModuleList(
            [nn.Linear(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):

    def __init__(self, n_mel_channels, postnet_embedding_dim,
            postnet_kernel_size, postnet_n_convolutions):

        super(Postnet, self).__init__()
        self.convolutions = nn.ModuleList()

        self.convolutions.append(
            nn.Sequential(
                nn.Conv1d(n_mel_channels, postnet_embedding_dim,
                          kernel_size=postnet_kernel_size, stride=1,
                          padding=postnet_kernel_size // 2),
                          dilation=1),
                nn.BatchNorm1d(postnet_embedding_dim))
        )

        for i in range(1, postnet_n_convolutions - 1):
            self.convolutions.append(
                nn.Sequential(
                    nn.Comv1d(postnet_embedding_dim,
                              postnet_embedding_dim,
                              postnet_kernel_size, stride=1,
                              padding=postnet_kernel_size // 2,
                              dilation=1),
                    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=postnet_kernel_size // 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

# Tacotron注意力机制
class LocationLayer(nn.Module):
    def __init__(self, attention_n_filters, attention_kernel_size,
                 attention_dim):
        super(LocationLayer, self).__init__()
        padding = attention_kernel_size // 2
        self.location_conv = nn.Conv2d(2, attention_n_filters,
                                      kernel_size=attention_kernel_size,
                                      padding=padding, bias=False, stride=1,
                                      dilation=1)
        self.location_dense = nn.Linear(attention_n_filters, attention_dim,
                                        bias=False)

    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 = nn.Linear(attention_rnn_dim,
                                     attention_dim,bias=False)
        self.memory_layer = nn.Linear(embedding_dim, 
                                      attention_dim, bias=False)

        self.v = nn.Linear(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):

        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):

        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

# Tacotron解码器
class Decoder(nn.Module):
    def __init__(self, n_mel_channels, n_frames_per_step,
        encoder_embedding_dim, attention_rnn_dim,
        decoder_rnn_dim, prenet_dim, max_decoder_steps,
        gate_threshold, p_attention_dropout,
        attention_dim, attention_location_n_filters,
        attention_location_kernel_size, p_decoder_dropout):

        super(Decoder, self).__init__()

        # 将输入参数保存到类的属性中
        # ... （此处省略保存输入参数的代码）
        self.prenet = Prenet(
            n_mel_channels * n_frames_per_step,
            [prenet_dim, prenet_dim])

        self.attention_rnn = nn.LSTMCell(
            prenet_dim + encoder_embedding_dim,
            attention_rnn_dim)

        self.attention_layer = Attention(
            attention_rnn_dim, encoder_embedding_dim,
            attention_dim, attention_location_n_filters,
            attention_location_kernel_size)

        self.decoder_rnn = nn.LSTMCell(
            attention_rnn_dim + encoder_embedding_dim,
            decoder_rnn_dim, 1)

        self.linear_projection = nn.Linear(
            decoder_rnn_dim + encoder_embedding_dim,
            n_mel_channels * n_frames_per_step)

        self.gate_layer = nn.Linear(
            decoder_rnn_dim + encoder_embedding_dim, 1,
            bias=True)

    def decode(self, decoder_input):
        # 输入解码器的梅尔过滤器特征，进行注意力机制的计算和循环神经网络计算
        # 输出解码结果，即是否终止的预测和注意力的权重
        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