从代码层面理解Tacotron中CBHG模块
人工智能技术

从代码层面理解Tacotron中CBHG模块

CBHG模块大致由三个层组成：1D卷积层，Highway层和双向GRU层

本文涉及到的代码来自 espnet/espnet/nets/pytorch_backend/tacotron2/cbhg.py at master · espnet/espnet (github.com)

代码中的一些默认超参数：

conv_bank_layers=8,
conv_bank_chans=128,
conv_proj_filts=3,
conv_proj_chans=256,
highway_layers=4,
highway_units=128,
gru_units=256,

本文将代码拆分为两块：定义模型和前向传播计算，两段基本上是分别连续的

先放一张原论文中的图：

6252783D-1280-478E-B8EC-A120FF089356.png

1D卷积层

建议先阅读：关于1D卷积-samyyc

代码：

# define 1d convolution bank
self.conv_bank = torch.nn.ModuleList()
for k in range(1, self.conv_bank_layers + 1):
    if k % 2 != 0:
        padding = (k - 1) // 2
    else:
        padding = ((k - 1) // 2, (k - 1) // 2 + 1)
    self.conv_bank += [
        torch.nn.Sequential(
            torch.nn.ConstantPad1d(padding, 0.0),
            torch.nn.Conv1d(
                idim, self.conv_bank_chans, k, stride=1, padding=0, bias=True
            ),
            torch.nn.BatchNorm1d(self.conv_bank_chans),
            torch.nn.ReLU(),
        )
    ]

定义这里的1D卷积层：

首先遍历k从 1 到 conv_bank_layers（超参数，对应论文中的 K sets），然后填充padding（这里padding的选用0），保证维度相同

如果卷积核的宽度是奇数的话，在两边填充一样的宽度即可，偶数就在尾部多填一个

注意这里的填充保证了所有卷积层卷积过后的输出维度都是相同的

然后进入1D卷积：

torch.nn.Conv1d(idim, self.conv_bank_chans, k, stride=1, padding=0, bias=True),

idim指输入向量的维度（就是这里参数的in_channels），self.conv_bank_chans指输出的维度（是个超参数），k指卷积核大小（也就是宽度）

然后进行BN，再ReLU激活

这里的输入维度是 (batch, idim, Tmax) ，Tmax指当前一个batch内时间序列的最大长度

忽略掉batch，等于想象一个矩阵，有idim行，Tmax列，每一列都代表一个序列中的一个input，此时一维卷积横向扫过这个矩阵，这个卷积核的长度等于idim（和输入矩阵的长度相同），宽度由遍历的k决定。为了保证Tmax一致，所以填充padding，一个通道的输出就对应维度： (batch, 1, Tmax) ，忽略掉batch就是一个Tmax长的特征向量，然后再把多个通道叠在一起，最后输出维度 (batch, channels, Tmax)

在进行此操作后，我们复刻了论文中的K组卷积层

来看这里对应的前向传播代码：


"""Calculate forward propagation.

Args:
    xs (Tensor): Batch of the padded sequences of inputs (B, Tmax, idim).
    ilens (LongTensor): Batch of lengths of each input sequence (B,).

Return:
    Tensor: Batch of the padded sequence of outputs (B, Tmax, odim).
    LongTensor: Batch of lengths of each output sequence (B,).

"""

xs = xs.transpose(1, 2)  # (B, idim, Tmax)
convs = []
for k in range(self.conv_bank_layers):
    convs += [self.conv_bank[k](xs)]
convs = torch.cat(convs, dim=1)  # (B, #CH * #BANK, Tmax)

首先原始xs的维度是 (batch, Tmax, idim) Tmax指一个batch内时间序列的最大长度，idim指输入向量的维度

先把xs第二和第三个维度交换，方便我们传入卷积层，然后把所有的卷积核都应用上计算一遍，得到一个结果列表（convs）

这里每一个卷积的输出维度都是 (batch, channels, Tmax) ，再在第二个维度上拼接就变成 (batch, banks * channels, Tmax) ，banks指有多少个过滤器，也就是原论文中的 K，上面的conv_bank_layers超参数

接着往下看：

# define max pooling (need padding for one-side to keep same length)
self.max_pool = torch.nn.Sequential(
            torch.nn.ConstantPad1d((0, 1), 0.0), torch.nn.MaxPool1d(2, stride=1)
        )

定义最大池化层，先在尾部填一个0，再两个两个一组最大池化，步数为1，还是保证维度相同，根据原论文，这里的 stride=1 维持了时间分辨率

对应的前向传播代码：

convs = self.max_pool(convs)

过程：

原始的convs维度： (batch, banks * channels, Tmax)

尾部填1后：(batch, banks * channels, Tmax+1)

池化后又回到 (batch, banks * channels, Tmax)

往下看：

这里是投影层（projection layer，其实就是全连接层）的部分：

# define 1d convolution projection
self.projections = torch.nn.Sequential(
    torch.nn.Conv1d(
        self.conv_bank_chans * self.conv_bank_layers,
        self.conv_proj_chans,
        self.conv_proj_filts,
        stride=1,
        padding=(self.conv_proj_filts - 1) // 2,
        bias=True,
    ),
    torch.nn.BatchNorm1d(self.conv_proj_chans),
    torch.nn.ReLU(),
    torch.nn.Conv1d(
        self.conv_proj_chans,
        self.idim,
        self.conv_proj_filts,
        stride=1,
        padding=(self.conv_proj_filts - 1) // 2,
        bias=True,
    ),
    torch.nn.BatchNorm1d(self.idim),
)

定义了两个卷积，第一个把通道数从 banks * channels 变为 conv_proj_chans，BN，激活，再卷积回idim，再BN

通过这一通操作后，维度回到 (batch, idim, Tmax)

这一部分投影层的作用其实就是控制维度

再来看前向传播代码：

convs = self.projections(convs).transpose(1, 2)  # (B, Tmax, idim)
xs = xs.transpose(1, 2) + convs

这边进行了transpose，把维度变为原来的 (batch, Tmax, idim) ，然后进行残差连接


HighwayNet层

这一部分主要介绍HighwayNet是什么

HighwayNet最主要的功能是跳过没有用的层，加快信息传递。

根据HighwayNet的原论文，作者将传统普通神经网络的非线性变换描述成： y =H(x, W_H)，其中H抽象的代表非线性变换， W_H代表变换中的参数

HighwayNet添加了两个非线性变换： T(x, W_T) 和 C(x, W_C) , T 被称为transform gate，C被称为carry gate，对于这两个名称的解释可以从下面这个式子体现：

y = H(x, W_H) \cdot T(x, W_T) + x \cdot C(x, W_C)

可以看到，transform gate控制了这个非线性变换层原来的输出，而carry gate控制了原来的输入，换句话说，transform gate控制有多少变换应该被输出，carry gate控制有多少输入应该被直接携带到输出

在原论文中，作者把 C 设为 1-T ，原式也就变成：

y = H(x, W_H) \cdot T(x, W_T) + x \cdot (1 - T(x, W_T))

再来看一下代码实现：

class HighwayNet(torch.nn.Module):
    """Highway Network module.

    This is a module of Highway Network introduced in `Highway Networks`_.

    .. _`Highway Networks`: https://arxiv.org/abs/1505.00387

    """

    def __init__(self, idim):
        """Initialize Highway Network module.

        Args:
            idim (int): Dimension of the inputs.

        """
        super(HighwayNet, self).__init__()
        self.idim = idim
        self.projection = torch.nn.Sequential(
            torch.nn.Linear(idim, idim), torch.nn.ReLU()
        )
        self.gate = torch.nn.Sequential(torch.nn.Linear(idim, idim), torch.nn.Sigmoid())

    def forward(self, x):
        """Calculate forward propagation.

        Args:
            x (Tensor): Batch of inputs (B, ..., idim).

        Returns:
            Tensor: Batch of outputs, which are the same shape as inputs (B, ..., idim).

        """
        proj = self.projection(x)
        gate = self.gate(x)
        return proj * gate + x * (1.0 - gate)

这一段还是比较好懂的，和上面的数学公式基本一致，维度没有变化

看一下CBHG中的代码：

# define highway network
self.highways = torch.nn.ModuleList()
self.highways += [torch.nn.Linear(idim, self.highway_units)]
for _ in range(self.highway_layers):
    self.highways += [HighwayNet(self.highway_units)]

第一层先是一个线性层，把维度从idim转成highwaynet网络的维度，然后重复highway_layers个highwaynet

前向传播的代码：

# + 1 for dimension adjustment layer
for i in range(self.highway_layers + 1):
    xs = self.highways[i](xs)

+1是为了把第一个转换维度的线性层也算进去

现在输出的维度： (batch, Tmax, HighwayUnits)

注意到默认参数里的 highway_units 和 conv_bank_chans 是一样的，所以一般情况下又回到了1D卷积层后的维度

双向GRU层

看一下代码里的定义：

self.gru = torch.nn.GRU(
    self.highway_units,
    gru_units // 2,
    num_layers=1,
    batch_first=True,
    bidirectional=True,
)

一个普通的双向GRU，隐藏层大小就是前面HighwayNet隐藏层的大小

注意到默认参数里 gru_units 是 highway_units 的两倍，这里又整除了2（为了双向），所以xs可以直接输入进去

前向传播：

# total_length needs for DataParallel
# (see https://github.com/pytorch/pytorch/pull/6327)
total_length = xs.size(1)
if not isinstance(ilens, torch.Tensor):
    ilens = torch.tensor(ilens)
xs = pack_padded_sequence(xs, ilens.cpu(), batch_first=True)
self.gru.flatten_parameters()
xs, _ = self.gru(xs)
xs, ilens = pad_packed_sequence(xs, batch_first=True, total_length=total_length)

# revert sorting by length
xs, ilens = self._revert_sort_by_length(xs, ilens, sort_idx)

这里的 flatten_parameters() 大概的用途是提高内存利用率和效率，使gru的内存变成连续的块

压缩原来的xs，并且输入到gru里，这里pack_padded_sequence就是用来压缩的，因为我们的xs输入的时候就已经被填充过了

然后再根据长度反向排序

全连接层

最后经过一层全连接层，转换维度到odim

# define final projection
self.output = torch.nn.Linear(gru_units, odim, bias=True)
xs = self.output(xs)  # (B, Tmax, odim)

return xs, ilens

最终输出维度: (batch, Tmax, odim)

tacotron
LICENSED UNDER CC BY-NC-SA 4.0
Comment