modules.py

import torch
import torch.nn as nn
import torch.nn.functional as F

from collections import OrderedDict
import numpy as np
import copy
import math

import hparams as hp
import utils

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def get_sinusoid_encoding_table(n_position, d_hid, padding_idx=None):
    ''' Sinusoid position encoding table '''

    def cal_angle(position, hid_idx):
        return position / np.power(10000, 2 * (hid_idx // 2) / d_hid)

    def get_posi_angle_vec(position):
        return [cal_angle(position, hid_j) for hid_j in range(d_hid)]

    sinusoid_table = np.array([get_posi_angle_vec(pos_i)
                               for pos_i in range(n_position)])

    sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
    sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

    if padding_idx is not None:
        # zero vector for padding dimension
        sinusoid_table[padding_idx] = 0.

    return torch.FloatTensor(sinusoid_table)


def clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])


class LengthRegulator(nn.Module):
    """ Length Regulator """

    def __init__(self):
        super(LengthRegulator, self).__init__()
        self.duration_predictor = DurationPredictor()

    def LR(self, x, duration_predictor_output, alpha=1.0, mel_max_length=None):
        output = list()

        for batch, expand_target in zip(x, duration_predictor_output):
            output.append(self.expand(batch, expand_target, alpha))

        if mel_max_length:
            output = utils.pad(output, mel_max_length)
        else:
            output = utils.pad(output)

        return output

    def expand(self, batch, predicted, alpha):
        out = list()

        for i, vec in enumerate(batch):
            expand_size = predicted[i].item()
            out.append(vec.expand(int(expand_size*alpha), -1))
        out = torch.cat(out, 0)

        return out

    def rounding(self, num):
        if num - int(num) >= 0.5:
            return int(num) + 1
        else:
            return int(num)

    def forward(self, x, alpha=1.0, target=None, mel_max_length=None):
        duration_predictor_output = self.duration_predictor(x)

        if self.training:
            output = self.LR(x, target, mel_max_length=mel_max_length)
            return output, duration_predictor_output
        else:
            for idx, ele in enumerate(duration_predictor_output[0]):
                duration_predictor_output[0][idx] = self.rounding(ele)
            output = self.LR(x, duration_predictor_output, alpha)
            mel_pos = torch.stack(
                [torch.Tensor([i+1 for i in range(output.size(1))])]).long().to(device)

            return output, mel_pos


# class DurationPredictor(nn.Module):
#     """ Duration Predictor """

#     def __init__(self):
#         super(DurationPredictor, self).__init__()

#         self.linear_layer_1 = Linear(hp.d_model, hp.d_model, w_init='relu')
#         self.relu_1 = nn.ReLU()
#         self.linear_layer_2 = Linear(hp.d_model, 1, w_init='relu')
#         self.relu_2 = nn.ReLU()

#     def forward(self, x):
#         out = self.linear_layer_1(x)
#         out = self.relu_1(out)
#         out = self.linear_layer_2(out)
#         out = self.relu_2(out)

#         out = out.squeeze()
#         if not self.training:
#             out = out.unsqueeze(0)

#         return out


# class LengthRegulator(nn.Module):
#     """ Length Regulator """

#     def __init__(self):
#         super(LengthRegulator, self).__init__()
#         self.duration_predictor = DurationPredictor()

#     def LR(self, encoder_output, duration_predictor_output, alpha, mel_max_length=None):
#         output = list()
#         for i in range(encoder_output.size(0)):
#             output.append(self.expand(
#                 encoder_output[i], duration_predictor_output[i], alpha))

#         if mel_max_length:
#             output, dec_pos = self.pad(output, mel_max_length)
#         else:
#             output, dec_pos = self.pad(output)

#         return output, dec_pos

#     def expand(self, one_batch, predicted, alpha):
#         out = list()
#         pad_length = list()

#         for ele in predicted:
#             pad_length.append(self.rounding(ele.data*alpha))

#         for i, ele in enumerate(one_batch):
#             out.append(ele.repeat(pad_length[i]+1, 1))
#         out = torch.cat(out, 0)

#         return out

#     def rounding(self, num):
#         if num - int(num) >= 0.5:
#             return int(num) + 1
#         else:
#             return int(num)

#     def pad(self, input_ele, mel_max_length=None):
#         if mel_max_length:
#             out_list = list()
#             max_len = mel_max_length

#             if input_ele[0].is_cuda:
#                 pos = torch.Tensor([i+1 for i in range(max_len)]
#                                    ).repeat(len(input_ele), 1).long().cuda()
#             else:
#                 pos = torch.Tensor([i+1 for i in range(max_len)]
#                                    ).repeat(len(input_ele), 1).long()

#             for i, batch in enumerate(input_ele):
#                 one_batch_padded = F.pad(
#                     batch, (0, 0, 0, max_len-batch.size(0)), "constant", 0.0)
#                 out_list.append(one_batch_padded)

#                 for ind in range(max_len-batch.size(0)):
#                     pos[i][batch.size(0)+ind] = 0

#             out_padded = torch.stack(out_list)
#             pos = pos.long()

#             return out_padded, pos
#         else:
#             out_list = list()
#             max_len = max([input_ele[i].size(0)
#                            for i in range(len(input_ele))])

#             if input_ele[0].is_cuda:
#                 pos = torch.Tensor([i+1 for i in range(max_len)]
#                                    ).repeat(len(input_ele), 1).long().cuda()
#             else:
#                 pos = torch.Tensor([i+1 for i in range(max_len)]
#                                    ).repeat(len(input_ele), 1).long()

#             for i, batch in enumerate(input_ele):
#                 one_batch_padded = F.pad(
#                     batch, (0, 0, 0, max_len-batch.size(0)), "constant", 0.0)
#                 out_list.append(one_batch_padded)

#                 for ind in range(max_len-batch.size(0)):
#                     pos[i][batch.size(0)+ind] = 0

#             out_padded = torch.stack(out_list)
#             pos = pos.long()

#             return out_padded, pos

#     def forward(self, encoder_output, target=None, alpha=1.0, mel_max_length=None):
#         duration_predictor_output = self.duration_predictor(encoder_output)

#         if self.training:
#             print(target)
#             output, decoder_pos = self.LR(
#                 encoder_output, target, alpha, mel_max_length)

#             return output, decoder_pos, duration_predictor_output
#         else:
#             # duration_predictor_output = torch.exp(duration_predictor_output)
#             # duration_predictor_output = duration_predictor_output - 1
#             output, decoder_pos = self.LR(
#                 encoder_output, duration_predictor_output, alpha)

#             return output, decoder_pos


class DurationPredictor(nn.Module):
    """ Duration Predictor """

    def __init__(self):
        super(DurationPredictor, self).__init__()

        self.input_size = hp.d_model
        self.filter_size = hp.duration_predictor_filter_size
        self.kernel = hp.duration_predictor_kernel_size
        self.conv_output_size = hp.duration_predictor_filter_size
        self.dropout = hp.dropout

        self.conv_layer = nn.Sequential(OrderedDict([
            ("conv1d_1", Conv(self.input_size,
                              self.filter_size,
                              kernel_size=self.kernel,
                              padding=1)),
            ("layer_norm_1", nn.LayerNorm(self.filter_size)),
            ("relu_1", nn.ReLU()),
            ("dropout_1", nn.Dropout(self.dropout)),
            ("conv1d_2", Conv(self.filter_size,
                              self.filter_size,
                              kernel_size=self.kernel,
                              padding=1)),
            ("layer_norm_2", nn.LayerNorm(self.filter_size)),
            ("relu_2", nn.ReLU()),
            ("dropout_2", nn.Dropout(self.dropout))
        ]))

        self.linear_layer = Linear(self.conv_output_size, 1)
        self.relu = nn.ReLU()

    def forward(self, encoder_output):
        out = self.conv_layer(encoder_output)
        out = self.linear_layer(out)

        out = self.relu(out)

        out = out.squeeze()

        if not self.training:
            out = out.unsqueeze(0)

        return out


class Conv(nn.Module):
    """
    Convolution Module
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=1,
                 stride=1,
                 padding=0,
                 dilation=1,
                 bias=True,
                 w_init='linear'):
        """
        :param in_channels: dimension of input
        :param out_channels: dimension of output
        :param kernel_size: size of kernel
        :param stride: size of stride
        :param padding: size of padding
        :param dilation: dilation rate
        :param bias: boolean. if True, bias is included.
        :param w_init: str. weight inits with xavier initialization.
        """
        super(Conv, self).__init__()

        self.conv = nn.Conv1d(in_channels,
                              out_channels,
                              kernel_size=kernel_size,
                              stride=stride,
                              padding=padding,
                              dilation=dilation,
                              bias=bias)

        nn.init.xavier_uniform_(
            self.conv.weight, gain=nn.init.calculate_gain(w_init))

    def forward(self, x):
        x = x.contiguous().transpose(1, 2)
        x = self.conv(x)
        x = x.contiguous().transpose(1, 2)

        return x


class Linear(nn.Module):
    """
    Linear Module
    """

    def __init__(self, in_dim, out_dim, bias=True, w_init='linear'):
        """
        :param in_dim: dimension of input
        :param out_dim: dimension of output
        :param bias: boolean. if True, bias is included.
        :param w_init: str. weight inits with xavier initialization.
        """
        super(Linear, self).__init__()
        self.linear_layer = nn.Linear(in_dim, out_dim, bias=bias)

        nn.init.xavier_uniform_(
            self.linear_layer.weight,
            gain=nn.init.calculate_gain(w_init))

    def forward(self, x):
        return self.linear_layer(x)


class FFN(nn.Module):
    """
    Positionwise Feed-Forward Network
    """

    def __init__(self, num_hidden):
        """
        :param num_hidden: dimension of hidden 
        """
        super(FFN, self).__init__()
        self.w_1 = Conv(num_hidden, num_hidden * 4,
                        kernel_size=3, padding=1, w_init='relu')
        self.w_2 = Conv(num_hidden * 4, num_hidden, kernel_size=3, padding=1)
        self.dropout = nn.Dropout(p=0.1)
        self.layer_norm = nn.LayerNorm(num_hidden)

    def forward(self, input_):
        # FFN Network
        x = input_
        x = self.w_2(torch.relu(self.w_1(x)))

        # residual connection
        x = x + input_

        # dropout
        x = self.dropout(x)

        # layer normalization
        x = self.layer_norm(x)

        return x


class MultiheadAttention(nn.Module):
    """
    Multihead attention mechanism (dot attention)
    """

    def __init__(self, num_hidden_k):
        """
        :param num_hidden_k: dimension of hidden 
        """
        super(MultiheadAttention, self).__init__()

        self.num_hidden_k = num_hidden_k
        self.attn_dropout = nn.Dropout(p=0.1)

    def forward(self, key, value, query, mask=None, query_mask=None):
        # Get attention score
        attn = torch.bmm(query, key.transpose(1, 2))
        attn = attn / math.sqrt(self.num_hidden_k)

        # Masking to ignore padding (key side)
        if mask is not None:
            attn = attn.masked_fill(mask, -2 ** 32 + 1)
            attn = torch.softmax(attn, dim=-1)
        else:
            attn = torch.softmax(attn, dim=-1)

        # Masking to ignore padding (query side)
        if query_mask is not None:
            attn = attn * query_mask

        # Dropout
        attn = self.attn_dropout(attn)

        # Get Context Vector
        result = torch.bmm(attn, value)

        return result, attn


class Attention(nn.Module):
    """
    Attention Network
    """

    def __init__(self, num_hidden, h=2):
        """
        :param num_hidden: dimension of hidden
        :param h: num of heads 
        """
        super(Attention, self).__init__()

        self.num_hidden = num_hidden
        self.num_hidden_per_attn = num_hidden // h
        self.h = h

        self.key = Linear(num_hidden, num_hidden, bias=False)
        self.value = Linear(num_hidden, num_hidden, bias=False)
        self.query = Linear(num_hidden, num_hidden, bias=False)

        self.multihead = MultiheadAttention(self.num_hidden_per_attn)

        self.residual_dropout = nn.Dropout(p=0.1)

        self.final_linear = Linear(num_hidden * 2, num_hidden)

        self.layer_norm_1 = nn.LayerNorm(num_hidden)

    def forward(self, memory, decoder_input, mask=None, query_mask=None):

        batch_size = memory.size(0)
        seq_k = memory.size(1)
        seq_q = decoder_input.size(1)

        # Repeat masks h times
        if query_mask is not None:
            query_mask = query_mask.unsqueeze(-1).repeat(1, 1, seq_k)
            query_mask = query_mask.repeat(self.h, 1, 1)
        if mask is not None:
            mask = mask.repeat(self.h, 1, 1)

        # Make multihead
        key = self.key(memory).view(batch_size,
                                    seq_k,
                                    self.h,
                                    self.num_hidden_per_attn)
        value = self.value(memory).view(batch_size,
                                        seq_k,
                                        self.h,
                                        self.num_hidden_per_attn)
        query = self.query(decoder_input).view(batch_size,
                                               seq_q,
                                               self.h,
                                               self.num_hidden_per_attn)

        key = key.permute(2, 0, 1, 3).contiguous().view(-1,
                                                        seq_k,
                                                        self.num_hidden_per_attn)
        value = value.permute(2, 0, 1, 3).contiguous().view(-1,
                                                            seq_k,
                                                            self.num_hidden_per_attn)
        query = query.permute(2, 0, 1, 3).contiguous().view(-1,
                                                            seq_q,
                                                            self.num_hidden_per_attn)

        # Get context vector
        result, attns = self.multihead(
            key, value, query, mask=mask, query_mask=query_mask)

        # Concatenate all multihead context vector
        result = result.view(self.h, batch_size, seq_q,
                             self.num_hidden_per_attn)
        result = result.permute(1, 2, 0, 3).contiguous().view(
            batch_size, seq_q, -1)

        # Concatenate context vector with input (most important)
        result = torch.cat([decoder_input, result], dim=-1)

        # Final linear
        result = self.final_linear(result)

        # Residual dropout & connection
        result = self.residual_dropout(result)
        result = result + decoder_input

        # Layer normalization
        result = self.layer_norm_1(result)

        return result, attns


class FFTBlock(torch.nn.Module):
    """FFT Block"""

    def __init__(self,
                 d_model,
                 n_head=hp.Head):
        super(FFTBlock, self).__init__()
        self.slf_attn = clones(Attention(d_model), hp.N)
        self.pos_ffn = clones(FFN(d_model), hp.N)

        self.pos_emb = nn.Embedding.from_pretrained(get_sinusoid_encoding_table(1024,
                                                                                d_model,
                                                                                padding_idx=0), freeze=True)

    def forward(self, x, pos, return_attns=False):
        # Get character mask
        if self.training:
            c_mask = pos.ne(0).type(torch.float)
            mask = pos.eq(0).unsqueeze(1).repeat(1, x.size(1), 1)
        else:
            c_mask, mask = None, None

        # Get positional embedding, apply alpha and add
        pos = self.pos_emb(pos)
        x = x + pos

        # Attention encoder-encoder
        attns = list()
        for slf_attn, ffn in zip(self.slf_attn, self.pos_ffn):
            x, attn = slf_attn(x, x, mask=mask, query_mask=c_mask)
            x = ffn(x)
            attns.append(attn)

        return x, attns