使用 Wav2Vec2 进行语音识别

本教程展示了如何使用 wav2vec 2.0 的预训练模型进行语音识别 [论文]。

概述

语音识别的过程如下所示。

从音频波形中提取声学特征
逐帧估计声学特征的类别
根据类别概率序列生成假设

Torchaudio 提供了对预训练权重及相关信息（如预期的采样率和类别标签）的便捷访问。它们被打包在一起，并可通过 torchaudio.pipelines 模块获取。

准备工作

import torch
import torchaudio

print(torch.__version__)
print(torchaudio.__version__)

torch.random.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print(device)

2.6.0
2.6.0
cuda

import IPython
import matplotlib.pyplot as plt
from torchaudio.utils import download_asset

SPEECH_FILE = download_asset("tutorial-assets/Lab41-SRI-VOiCES-src-sp0307-ch127535-sg0042.wav")

  0%|          | 0.00/106k [00:00<?, ?B/s]
100%|##########| 106k/106k [00:00<00:00, 62.6MB/s]

创建管道

首先，我们将创建一个 Wav2Vec2 模型，用于执行特征提取和分类。

在 torchaudio 中，有两种类型的 Wav2Vec2 预训练权重：一种是针对 ASR 任务进行微调的，另一种是未进行微调的。

Wav2Vec2（以及 HuBERT）模型是以自监督的方式进行训练的。它们首先仅使用音频进行表示学习，然后使用额外的标签针对特定任务进行微调。

未进行微调的预训练权重也可以用于其他下游任务的微调，但本教程不涉及这部分内容。

我们将在这里使用 torchaudio.pipelines.WAV2VEC2_ASR_BASE_960H。

在 torchaudio.pipelines 中提供了多个预训练模型。请查阅文档以了解它们的训练细节。

该 bundle 对象提供了实例化模型和其他信息的接口。采样率和类别标签如下所示。

bundle = torchaudio.pipelines.WAV2VEC2_ASR_BASE_960H

print("Sample Rate:", bundle.sample_rate)

print("Labels:", bundle.get_labels())

Sample Rate: 16000
Labels: ('-', '|', 'E', 'T', 'A', 'O', 'N', 'I', 'H', 'S', 'R', 'D', 'L', 'U', 'M', 'W', 'C', 'F', 'G', 'Y', 'P', 'B', 'V', 'K', "'", 'X', 'J', 'Q', 'Z')

模型可以按以下方式构建。此过程将自动获取预训练权重并将其加载到模型中。

model = bundle.get_model().to(device)

print(model.__class__)

Downloading: "https://download.pytorch.org/torchaudio/models/wav2vec2_fairseq_base_ls960_asr_ls960.pth" to /root/.cache/torch/hub/checkpoints/wav2vec2_fairseq_base_ls960_asr_ls960.pth

  0%|          | 0.00/360M [00:00<?, ?B/s]
 16%|#5        | 56.8M/360M [00:00<00:00, 594MB/s]
 32%|###1      | 114M/360M [00:00<00:00, 595MB/s]
 47%|####7     | 170M/360M [00:00<00:00, 595MB/s]
 63%|######3   | 227M/360M [00:00<00:00, 585MB/s]
 79%|#######8  | 284M/360M [00:00<00:00, 589MB/s]
 95%|#########4| 341M/360M [00:00<00:00, 591MB/s]
100%|##########| 360M/360M [00:00<00:00, 590MB/s]
<class 'torchaudio.models.wav2vec2.model.Wav2Vec2Model'>

加载数据

我们将使用来自 VOiCES 数据集的语音数据，该数据集遵循 Creative Commons BY 4.0 许可协议。

IPython.display.Audio(SPEECH_FILE)

要加载数据，我们使用 torchaudio.load()。

如果采样率与管道期望的不同，我们可以使用 torchaudio.functional.resample() 进行重采样。

torchaudio.functional.resample() 也可以处理 CUDA 张量。

如果需要在同一组采样率上进行多次重采样，使用 torchaudio.transforms.Resample 可能会提高性能。

waveform, sample_rate = torchaudio.load(SPEECH_FILE)
waveform = waveform.to(device)

if sample_rate != bundle.sample_rate:
    waveform = torchaudio.functional.resample(waveform, sample_rate, bundle.sample_rate)

提取声学特征

下一步是从音频中提取声学特征。

针对ASR任务进行微调的Wav2Vec2模型可以一步完成特征提取和分类，但为了本教程的完整性，我们在此也展示了如何进行特征提取。

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

返回的特征是一个张量列表。每个张量都是 transformer 层的输出。

fig, ax = plt.subplots(len(features), 1, figsize=(16, 4.3 * len(features)))
for i, feats in enumerate(features):
    ax[i].imshow(feats[0].cpu(), interpolation="nearest")
    ax[i].set_title(f"Feature from transformer layer {i+1}")
    ax[i].set_xlabel("Feature dimension")
    ax[i].set_ylabel("Frame (time-axis)")
fig.tight_layout()

特征分类

一旦提取了音频特征，下一步就是将它们分类到一组类别中。

Wav2Vec2 模型提供了在一步中执行特征提取和分类的方法。

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

输出以 logits 的形式呈现，而不是概率形式。

让我们将其可视化。

plt.imshow(emission[0].cpu().T, interpolation="nearest")
plt.title("Classification result")
plt.xlabel("Frame (time-axis)")
plt.ylabel("Class")
plt.tight_layout()
print("Class labels:", bundle.get_labels())

Classification result

Class labels: ('-', '|', 'E', 'T', 'A', 'O', 'N', 'I', 'H', 'S', 'R', 'D', 'L', 'U', 'M', 'W', 'C', 'F', 'G', 'Y', 'P', 'B', 'V', 'K', "'", 'X', 'J', 'Q', 'Z')

我们可以看到，在整个时间线上，某些标签的指示非常强烈。

生成转录文本

从标签概率序列中，我们现在希望生成转录文本。生成假设的过程通常被称为“解码”。

解码比简单的分类更为复杂，因为在某个时间步的解码可能会受到周围观测值的影响。

例如，以单词 night 和 knight 为例。即使它们的先验概率分布不同（在典型的对话中，night 的出现频率远高于 knight），为了准确生成包含 knight 的转录文本，例如 a knight with a sword，解码过程必须推迟最终决策，直到看到足够的上下文。

已经提出了许多解码技术，这些技术需要外部资源，例如单词词典和语言模型。

在本教程中，为了简单起见，我们将执行贪心解码，它不依赖于这些外部组件，而是简单地在每个时间步选择最佳假设。因此，上下文信息不会被使用，并且只能生成一个转录文本。

我们从定义贪心解码算法开始。

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:
"""Given a sequence emission over labels, get the best path string
        Args:
          emission (Tensor): Logit tensors. Shape `[num_seq, num_label]`.

        Returns:
          str: The resulting transcript
        """
        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)
IPython.display.Audio(SPEECH_FILE)

I|HAD|THAT|CURIOSITY|BESIDE|ME|AT|THIS|MOMENT|

ASR 模型使用一种称为连接时序分类（CTC）的损失函数进行微调。CTC 损失的详细信息在这里进行了说明。在 CTC 中，空白标记（ϵ）是一个特殊标记，表示前一个符号的重复。在解码过程中，这些标记会被简单地忽略。

结论

在本教程中，我们探讨了如何使用 Wav2Vec2ASRBundle 进行声学特征提取和语音识别。构建模型并获取输出仅需两行代码。

model = torchaudio.pipelines.WAV2VEC2_ASR_BASE_960H.get_model()
emission = model(waveforms, ...)

下载 Python 源代码: speech_recognition_pipeline_tutorial.py

下载 Jupyter 笔记本: speech_recognition_pipeline_tutorial.ipynb