前言

wav2vec 2.0是由Facebook 人工智能实验室（Facebook AI Research Paris, FARP）与2020年发表在Advances in Neural Information Processing Systems 期刊的语音领域的论文。该论文提出的无监督语音表示框架：其核心就相当于是一种语音特征的提取器，这种提取器可以提取很多的音频通用特征，而这些通用特征是可以运用在语音领域中微调后的下游任务如：语音识别，声纹识别，多轮对话。。。

模型架构

这里我以Pytorch中的wav2vec2-base英文模型和腾讯游戏知几AI团队与西工大ASLP组联合发布的中文预训练模型为例子进行讲解。wav2vec 2.0 模型的整体架构图如下：

wav2vec 2.0 模型整体架构图

其中的seq_num 表示帧长，Conv1d表示一维卷积网络；Fp32GroupNorm表示实际上是Pytorch中的GroupNorm，不过fairseq里面的层名字设置是这样的；PQ就是乘积量化；FC就是全连接层；LN就是LayerNorm；Self-Attention表示的是注意力头，每一层有8个；FFN前馈神经网络实际上也是一种全连接网络其内置维度是3072；PAD-MASK 就是掩码层，估计就是0 值掩码，掩码策略：选6.5%的时间步长作为开始序号，每个序号后面的十个时间步长被掩码。

中文&英文预训练模型差异

虽然整体架构就如上图所示，但是实际上因为调参和使用神经网络框架不同的原因会有些许区别。比如：Dropout层使用数量和LayerNorm层使用顺序等。举个特征提取层的例子，以下是中文预训练时的卷积特征提取层的部分模型：

Wav2Vec2Model(
  (feature_extractor): ConvFeatureExtractionModel(
    (conv_layers): ModuleList(
      (0): Sequential(
        (0): Conv1d(1, 512, kernel_size=(10,), stride=(5,), bias=False)
        (1): Dropout(p=0.0, inplace=False)
        (2): Fp32GroupNorm(512, 512, eps=1e-05, affine=True)
        (3): GELU()
      )
      (1): Sequential(
        (0): Conv1d(512, 512, kernel_size=(3,), stride=(2,), bias=False)
        (1): Dropout(p=0.0, inplace=False)
        (2): GELU()
      )

以下是英文预训练时的卷积特征提取层的部分模型：

Wav2Vec2Model(
  (feature_extractor): FeatureExtractor(
    (conv_layers): ModuleList(
      (0): ConvLayerBlock(
        (layer_norm): GroupNorm(512, 512, eps=1e-05, affine=True)
        (conv): Conv1d(1, 512, kernel_size=(10,), stride=(5,), bias=False)
      )
      (1): ConvLayerBlock(
        (conv): Conv1d(512, 512, kernel_size=(3,), stride=(2,), bias=False)
      )

很显然可以看到英文的预训练的GELU激活函数没有使用，但是实际上作者在论文中提到过他是用了GELU函数的。我认为这是AI神经网络的工程师们在实现不同语音预训练的时候不同的调参策略。

代码解析

首先是Pytorch代码解析，由于暂时还没做好中文预训练模型下的自动语音识别（Automatic Speech Recognition,ASR）任务，暂时先谈谈英文的。

1. pytorch下的ASR任务代码demo讲解

首先是代码的部分的展示，如下面的代码块所示：

import torchaudio
import torch
SPEECH_FILE = 'resources/En-1.wav'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
bundle = torchaudio.pipelines.WAV2VEC2_ASR_BASE_960H
model = bundle.get_model().to(device)  # 获取模型
print(model)
labels = bundle.get_labels()  # 获取字典标签
waveform, sample_rate = torchaudio.load(SPEECH_FILE)
waveform = waveform.to(device)

with torch.inference_mode():
  features, _ = model.extract_features(waveform)

with torch.inference_mode():
  emission, _ = model(waveform)


class GreedyCTCDecoder(torch.nn.Module):
  def __init__(self, labels, blank=0):
    super().__init__()
    self.labels = labels
    self.blank = blank

  def forward(self, emission: torch.Tensor) -> str:
    indices = torch.argmax(emission, dim=-1)  # [num_seq,]
    indices = torch.unique_consecutive(indices, dim=-1)
    indices = [i for i in indices if i != self.blank]
    return ''.join([self.labels[i] for i in indices])


decoder = GreedyCTCDecoder(labels=bundle.get_labels())
transcript = decoder(emission[0])
print(transcript)

其中先导入torch、torchaudio模块，torchaudio中的pipelines是为了方便执行某一项具体任务而讲训练过程中对音频的处理以及对模型推理之后的概率分布处理操作或其他后置处理封装在一起的API。features=[1,seq_num,768]是12层Transfomer-Encode特征提取之后得到的最终结果。GreedyCTCDecoder类是一个简单的CTC对齐的方法，他是用来对概率最大的标签进行匹配的得到想要的结果。

2. HuggingFace的ASR任务

HuggingFace代码部分的展示，如下面的代码块所示：

import soundfile as sf
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor


path_audio = 'resources/En-1.wav'
processor = Wav2Vec2Processor.from_pretrained("C:/Users/12046/.cache/huggingface/transformers")
model = Wav2Vec2ForCTC.from_pretrained("C:/Users/12046/.cache/huggingface/transformers")   # 用于ASR等，32维

audio_input, sample_rate = sf.read(path_audio)  # (31129,)
input_values = processor(audio_input, sampling_rate=sample_rate, return_tensors="pt").input_values  # torch.Size([1, 31129])

logits = model(input_values).logits     # torch.Size([1, 97, 32])
predicted_ids = torch.argmax(logits, dim=-1)    # torch.Size([1, 97])

transcription = processor.decode(predicted_ids[0])  # ASR的解码结果
print(transcription)

HuggingFace经常是用的就是transfomers这个库了，实际上这里的processor用的是和pytorch代码中的GreedyCTCDecoder一个思想，不过因为他里面没有pytorch中方便的pipelines所以不能直接进行解码操作，需要一个后置处理的类进行解码操作。

总结

本篇博客，对wav2vec2.0 的模型架构和代码进行了简单的讲解，希望看到这篇博客并且在进行语音方向学习的人有所收货lol。