当前位置：首页 > article >正文

VITS 源码解析2-模型概述

article 2024/11/15 15:03:56

VITs是文本到语音(Text-to-Speech, TTS)任务中最流行的技术之一，其实现思路是将文本语音信息融合到了HiFiGAN潜空间内, 通过文本控制HiFiGAN的生成器，输出含文本语义的声音。

VITs主要以GAN的方式训练, 其生成器G是SynthesizerTrn，判别器D是MPD。

VITS的判别器几乎和HiFiGAN一样，生成器则融合了文本、时序、声音三大类模型

1.文件概述

模型部分包含三个文件

attentons.py

注意机制就是transformer，在文本编码器中用到了，transformer的encoder。
modules.py

这个包含模型的一些基础结构(blocks), 比如Norm, Conv, Resblock等
models.py

这个是VITS核心的模型结构，前两个文件只是其基础。

2.attentions.py

就是标准的transformer实现，包括

Encoder
Decoder
MultiHeadAttention
FFN (Feed-Forward Network)

前两个是模型，后两个是模型的block:

1.MultiHeadAttention处理token之间的关系

2.FFN处理单个token的字符关系。

3.modules.py

3.1 Norm

LayerNorm

LayerNorm 类是一个层归一化（Layer Normalization）的自定义模块，用于对输入的特征进行归一化，旨在加速神经网络的训练并提升模型的稳定性。

1.层归一化是对一个层的所有神经元的激活值进行归一化处理，确保输出具有稳定的分布。

2.有学习的 gamma 和 beta 参数，用于学归一化后调整输出的尺度和偏移。

3.维度交换 是为了适应 PyTorch 的 layer_norm 函数的输入格式，归一化完毕后再将维度还原

3.2 卷积块

ConvReluNorm

ConvReluNorm 结合了1d卷积层、ReLU激活函数、Dropout和LayerNorm，用于处理序列数据。

在设计上，它通过多层卷积对输入进行逐步处理，并且通过添加归一化和激活函数来增强模型的表达能力，是模型的通用模块。

DDSConv

DDSConv 是结合膨胀卷积（dilation）和深度可分离卷积（groups），保持局部感受野的高效性，也通过膨胀卷积逐步扩大感受野。

1.使用LayerNorm进行归一化，确保各层之间的数值稳定性。

2.引入GELU激活函数和Dropout，进一步增强非线性表达能力并防止过拟合。

3.通过残差连接，有效地缓解了深层网络中的梯度消失问题，使模型可以更深层次地训练。

3.3 WaveNet

即WN，可用于生成waveforms

主要使用膨胀卷积、条件输入和残差跳跃连接。

3.4 残差块

ResBlock1
1. 包含两个独立的卷积模块（self.convs1 和 self.convs2），每个模块有三层1d卷积
2. self.convs1 的各层应用不同的膨胀率（1, 3, 5），这些膨胀率允许网络捕捉更大的感受野
3.self.convs2各层用统一膨胀率
ResBlock2
1. 仅有2个1d卷积，使用不同膨胀率

ResBlock1 提供了更多的卷积层和不同的膨胀率配置，以增强模型的表达能力，而 ResBlock2 提供了一个较简化的结构。

3.5 操作函数

这里用三个类分别实现了不同的正则化变换，用于完成流模型的生成和可逆变换:

Log：实现对数变换，用于数据的对数处理和行列式计算。
Flip：实现张量翻转操作，翻转操作不会改变行列式的绝对值（绝对值为 1）），其行列式的对数值为零。
ElementwiseAffine：实现逐元素的仿射变换，包括缩放和偏移，并计算相应的行列式的对数。

以下对每个类的详细解读：

Log类

Log 类实现了对输入张量 x 的对数变换。它的正向变换是取对数，逆变换是取指数。

1.正向变换计算 y 为 x 的对数值，并通过 torch.clamp_min 确保 x 的最小值为 1e-5，避免对数计算中的数值不稳定，返回变换后的张量 y 和 logdet。

2.逆向变换计算 x 为 y 的指数值，并乘以掩码 x_mask，以恢复原始输入，返回变换后的张量 x。

代码：

class Log(nn.Module):
	def forward(self, x, x_mask, reverse=False, **kwargs): 
		if not reverse: 
			y = torch.log(torch.clamp_min(x, 1e-5)) * x_mask 
			logdet = torch.sum(-y, [1, 2]) 
		return y,  logdet 

		else:
			 x = torch.exp(x) * x_mask return x

Flip

Flip 类实现了对输入张量 x 的翻转变换，值没有实际变化。

1.正向变换：
1. 使用 torch.flip 在指定的维度（这里是维度 1）对张量 x 进行翻转。
2. logdet 被设置为零张量，因为翻转操作的行列式是 1（对数行列式为 0）。
3. 返回翻转后的张量 x 和 logdet。

2.逆向变换：
	由于翻转操作是对称的，所以逆变换也仅是翻转操作本身，直接返回张量 x。

代码：

class Flip(nn.Module):
  def forward(self, x, *args, reverse=False, **kwargs):
    x = torch.flip(x, [1])
    if not reverse:
      logdet = torch.zeros(x.size(0)).to(dtype=x.dtype, device=x.device)
      return x, logdet
    else:
      return x

ElementwiseAffine

实现了一个逐元素的仿射变换，每个通道变换包括缩放和偏移操作。

1.正向变换：
1.对输入 x 应用仿射变换 y = m + exp(logs) * x，其中 m 和 logs 是可学习的参数。
2.应用掩码 x_mask，确保在变换时考虑掩码的影响。
3.计算 logdet，即 logs 与掩码的乘积的和，用于计算行列式的对数。
4.返回变换后的张量 y 和 logdet。

2.逆向变换：
	1.使用 m 和 logs 的参数来恢复原始输入 x，即 x = (x - m) * exp(-logs)。
	2.返回逆变换后的张量 x。

代码：

class ElementwiseAffine(nn.Module):
  def __init__(self, channels):
    super().__init__()
    self.channels = channels
    self.m = nn.Parameter(torch.zeros(channels,1))
    self.logs = nn.Parameter(torch.zeros(channels,1))

  def forward(self, x, x_mask, reverse=False, **kwargs):
    if not reverse:
      y = self.m + torch.exp(self.logs) * x
      y = y * x_mask
      logdet = torch.sum(self.logs * x_mask, [1,2])
      return y, logdet
    else:
      x = (x - self.m) * torch.exp(-self.logs) * x_mask
      return x

3.6 流模型

这部分实现了两个类，都是耦合层结构，负责对输入进行变换，及其可逆变换。

ResidualCouplingLayer 使用较为简单的线性仿射变换，而 ConvFlow 使用非线性的有理二次分段变换（Rational Quadratic Spline），适合处理更加复杂的分布调整。

ResidualCouplingLayer

基于残差耦合结构，输入 x 按照通道维度拆分为两部分 x0 和 x1，然后仅对 x1 进行变换，x0保持不变，

x1的变换参数是非线性的，依赖于x0。

过程：

1. x分割为 x0, x1, 

2. 通过卷积层 (self.pre) 对 x0 进行预处理, 得到h

3. 将h输入到一个 WN（WaveNet 样的网络）中，WN由多层膨胀卷积和门控激活单元组成，
用于生成 x1 的条件变换参数,  m 和 logs，可对 x1 进行平移和尺度变换。

5. 线性仿射变换:   正向变换  x1 = m + x1 *exp(logs),  逆向变换 x1 = (x1-m)*exp(-logs)

6. 前向变换会额外输出logdet，logs 是用于尺度变换的对数值， logdet 是对 logs 的加总，用于计算整体变换的行列式。

ConvFlow

1.类似 ResidualCouplingLayer，x分割为 x0, x1，x0输入pre层，得到h。

2. h继续输入深度分离膨胀卷积 (DDSConv) ，提取特征，参数为mask和speakers

3. h 继续输入到投影层 (self.proj) 生成与分段有理二次变换相关的多个参数，
包括 unnormalized_widths（未归一化的区间宽度）、
unnormalized_heights（区间高度）和 unnormalized_derivatives（导数）。

4. 使用from transforms 的 piecewise_rational_quadratic_transform 函数对 x1 进行分段
有理二次变换，得到最终的x1 和 logdet

5.logdet：在这个变换中，logdet 是通过有理二次变换的计算得出，它反映了变换过程中对数行列式的变化。

4.models.py

4.1 文本融合部分

StochasticDurationPredictor: Flow模型实现
DurationPredictor：标准解码器，1D卷积网络实现
TextEncoder: 用于将文本编码为潜空间向量。更准确地说，是将音素编码为潜向量。
ResidualCouplingBlock: Flow模型的block，实现逆计算(Inverse)
SynthesizerTrn: 将多个模型封装在一个类

4.2 HiFiGAN部分

PosteriorEncoder: HiFi-GAN部分的模型，这里叫后验编码器，用于将mel时频谱编码为潜向量z
Generator: HiFi-GAN部分的生成器，即解码器，输入潜向量z，输出对于音频waveforms
DiscriminatorP: 周期判别器

可检测音频信号中的周期性模式，周期性是许多自然音频（例如声音的频率分量）中的一个关键特征。

步骤：

* 将一维的音频信号转换为二维的表示形式（时间-周期的形式）。

* 利用卷积层逐步提取特征，并且通过每一层的不同卷积核大小和步幅（stride）来捕捉音频中的周期性特征。

* 使用不同的周期长度 period 来划分音频信号，这样可以检测到不同频率的周期性模式。

DiscriminatorS: 多尺度判别器

用于在多个时间尺度上对音频信号进行判别。判别器通过不同大小的卷积核和步幅来检测音频实现。

设计网络目的是捕捉音频中不同时间尺度的特征, 判断生成的音频是否自然。

步骤：

* 通过多层一维卷积操作，逐层提取音频信号中的特征。

* 使用多个不同尺度的卷积核来捕捉音频中从局部到全局的时间模式。

* 最终使用一个 1D 卷积层进行后处理，并通过 conv_post 得到打平的输出。

代码：

每一层卷积核的大小（如 15, 41）以及组卷积（groups=4, 16, 64）逐渐增加，可以帮助捕捉不同尺度的音频特征，提取越来越细致的音频特征。

self.convs = nn.ModuleList([
    norm_f(Conv1d(1, 16, 15, 1, padding=7)),
    norm_f(Conv1d(16, 64, 41, 4, groups=4, padding=20)),
    norm_f(Conv1d(64, 256, 41, 4, groups=16, padding=20)),
    norm_f(Conv1d(256, 1024, 41, 4, groups=64, padding=20)),
    norm_f(Conv1d(1024, 1024, 41, 4, groups=256, padding=20)),
    norm_f(Conv1d(1024, 1024, 5, 1, padding=2)),
])

self.conv_post = norm_f(Conv1d(1024, 1, 3, 1, padding=1))

MultiPeriodDiscriminator

多尺度判别器 (DiscriminatorS) ：专注于不同时间尺度上的特征，对生成音频的整体结构（如长时间的音高变化、节奏等)进行判断
周期判别器 (DiscriminatorP) ：可以捕捉音频中的周期性模式，尤其是声音的频率和音调特征。

模型结构:

periods = [2,3,5,7,11] # 选择不同的周期 

discs = [DiscriminatorS(use_spectral_norm=use_spectral_norm)] # 多尺度判别器 

discs = discs + [DiscriminatorP(i, use_spectral_norm=use_spectral_norm) for i in periods] # 添加多个周期判别器 

self.discriminators = nn.ModuleList(discs)