使用 LM 头微调 Wav2Vec2

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看

下载笔记本

查看 TF Hub 模型

在本笔记本中，我们将从 TFHub 加载预训练的 wav2vec2 模型，并通过在预训练模型的顶部追加语言建模头 (LM) 在 LibriSpeech 数据集 上对其进行微调。底层任务是构建一个用于 **自动语音识别** 的模型，即给定一些语音，模型应该能够将其转录成文本。

设置

在运行本笔记本之前，请确保您使用的是 GPU 运行时 (Runtime > Change runtime type > GPU)。以下单元格将安装 gsoc-wav2vec2 包及其依赖项。

pip3 install -q git+https://github.com/vasudevgupta7/gsoc-wav2vec2@main
sudo apt-get install -y libsndfile1-dev
pip3 install -q SoundFile

使用 TFHub 设置模型

我们将从导入一些库/模块开始。

import os

import tensorflow as tf
import tensorflow_hub as hub
from wav2vec2 import Wav2Vec2Config

config = Wav2Vec2Config()

print("TF version:", tf.__version__)

首先，我们将从 TFHub 下载我们的模型，并将我们的模型签名与 hub.KerasLayer 包装起来，以便能够像使用任何其他 Keras 层一样使用此模型。幸运的是，hub.KerasLayer 只需一行代码即可完成这两项操作。

pretrained_layer = hub.KerasLayer("https://tfhub.dev/vasudevgupta7/wav2vec2/1", trainable=True)

如果您有兴趣了解模型导出脚本，可以参考 此脚本。对象 pretrained_layer 是 Wav2Vec2Model 的冻结版本。这些预训练权重是使用 此脚本 从 HuggingFace PyTorch 预训练权重 转换而来的。

最初，wav2vec2 使用掩码语言建模方法进行预训练，其目标是识别掩码时间步的真实量化潜在语音表示。您可以在论文中阅读有关训练目标的更多信息 - wav2vec 2.0：一种用于语音表示的自监督学习框架。

现在，我们将定义一些常量和超参数，这些常量和超参数在接下来的几个单元格中将很有用。 AUDIO_MAXLEN 故意设置为 246000，因为模型签名只接受 246000 的静态序列长度。

AUDIO_MAXLEN = 246000
LABEL_MAXLEN = 256
BATCH_SIZE = 2

在以下单元格中，我们将使用 Keras 的函数式 API 包装 pretrained_layer 和一个密集层 (LM 头)。

inputs = tf.keras.Input(shape=(AUDIO_MAXLEN,))
hidden_states = pretrained_layer(inputs)
outputs = tf.keras.layers.Dense(config.vocab_size)(hidden_states)

model = tf.keras.Model(inputs=inputs, outputs=outputs)

密集层（如上定义）的输出维度为 vocab_size，因为我们希望在每个时间步预测词汇表中每个标记的概率。

设置训练状态

在 TensorFlow 中，模型权重仅在第一次调用 model.call 或 model.build 时构建，因此以下单元格将为我们构建模型权重。此外，我们将运行 model.summary() 来检查可训练参数的总数。

model(tf.random.uniform(shape=(BATCH_SIZE, AUDIO_MAXLEN)))
model.summary()

现在，我们需要定义 loss_fn 和优化器才能训练模型。以下单元格将为我们执行此操作。我们将使用 Adam 优化器，以简化操作。 CTCLoss 是一种常见的损失类型，用于输入子部分难以与输出子部分对齐的任务（如 ASR）。您可以从这篇精彩的 博客文章 中了解更多关于 CTC 损失的信息。

CTCLoss (来自 gsoc-wav2vec2 包) 接受 3 个参数：config、model_input_shape 和 division_factor。如果 division_factor=1，则损失将简单地进行累加，因此请根据需要传递 division_factor 以获取批次的平均值。

from wav2vec2 import CTCLoss

LEARNING_RATE = 5e-5

loss_fn = CTCLoss(config, (BATCH_SIZE, AUDIO_MAXLEN), division_factor=BATCH_SIZE)
optimizer = tf.keras.optimizers.Adam(LEARNING_RATE)

加载和预处理数据

现在，让我们从 官方网站 下载 LibriSpeech 数据集并进行设置。

wget https://www.openslr.org/resources/12/dev-clean.tar.gz -P ./data/train/
tar -xf ./data/train/dev-clean.tar.gz -C ./data/train/

ls ./data/train/

我们的数据集位于 LibriSpeech 目录中。让我们探索这些文件。

data_dir = "./data/train/LibriSpeech/dev-clean/2428/83705/"
all_files = os.listdir(data_dir)

flac_files = [f for f in all_files if f.endswith(".flac")]
txt_files = [f for f in all_files if f.endswith(".txt")]

print("Transcription files:", txt_files, "\nSound files:", flac_files)

好的，每个子目录都有许多 .flac 文件和一个 .txt 文件。该 .txt 文件包含该子目录中所有语音样本（即 .flac 文件）的文本转录。

我们可以按如下方式加载此文本数据

def read_txt_file(f):
  with open(f, "r") as f:
    samples = f.read().split("\n")
    samples = {s.split()[0]: " ".join(s.split()[1:]) for s in samples if len(s.split()) > 2}
  return samples

类似地，我们将定义一个函数，用于从 .flac 文件加载语音样本。

REQUIRED_SAMPLE_RATE 设置为 16000，因为 wav2vec2 是使用 16K 频率预训练的，建议在没有对数据分布进行重大更改的情况下进行微调，因为频率的原因。

import soundfile as sf

REQUIRED_SAMPLE_RATE = 16000

def read_flac_file(file_path):
  with open(file_path, "rb") as f:
      audio, sample_rate = sf.read(f)
  if sample_rate != REQUIRED_SAMPLE_RATE:
      raise ValueError(
          f"sample rate (={sample_rate}) of your files must be {REQUIRED_SAMPLE_RATE}"
      )
  file_id = os.path.split(file_path)[-1][:-len(".flac")]
  return {file_id: audio}

现在，我们将选择一些随机样本，并尝试对其进行可视化。

from IPython.display import Audio
import random

file_id = random.choice([f[:-len(".flac")] for f in flac_files])
flac_file_path, txt_file_path = os.path.join(data_dir, f"{file_id}.flac"), os.path.join(data_dir, "2428-83705.trans.txt")

print("Text Transcription:", read_txt_file(txt_file_path)[file_id], "\nAudio:")
Audio(filename=flac_file_path)

现在，我们将组合所有语音和文本样本，并将定义该目的的函数（在下一个单元格中）。

def fetch_sound_text_mapping(data_dir):
  all_files = os.listdir(data_dir)

  flac_files = [os.path.join(data_dir, f) for f in all_files if f.endswith(".flac")]
  txt_files = [os.path.join(data_dir, f) for f in all_files if f.endswith(".txt")]

  txt_samples = {}
  for f in txt_files:
    txt_samples.update(read_txt_file(f))

  speech_samples = {}
  for f in flac_files:
    speech_samples.update(read_flac_file(f))

  assert len(txt_samples) == len(speech_samples)

  samples = [(speech_samples[file_id], txt_samples[file_id]) for file_id in speech_samples.keys() if len(speech_samples[file_id]) < AUDIO_MAXLEN]
  return samples

现在是时候看看一些样本了...

samples = fetch_sound_text_mapping(data_dir)
samples[:5]

现在让我们预处理数据！！！

我们首先将使用 gsoc-wav2vec2 包定义分词器和处理器。然后，我们将进行非常简单的预处理。processor 将针对帧轴对原始语音进行归一化，而 tokenizer 将将我们的模型输出转换为字符串（使用定义的词汇表），并将负责删除特殊标记（取决于您的分词器配置）。

from wav2vec2 import Wav2Vec2Processor
tokenizer = Wav2Vec2Processor(is_tokenizer=True)
processor = Wav2Vec2Processor(is_tokenizer=False)

def preprocess_text(text):
  label = tokenizer(text)
  return tf.constant(label, dtype=tf.int32)

def preprocess_speech(audio):
  audio = tf.constant(audio, dtype=tf.float32)
  return processor(tf.transpose(audio))

现在，我们将定义 Python 生成器来调用我们在上面单元格中定义的预处理函数。

def inputs_generator():
  for speech, text in samples:
    yield preprocess_speech(speech), preprocess_text(text)

设置 tf.data.Dataset

以下单元格将使用其 .from_generator(...) 方法设置 tf.data.Dataset 对象。我们将使用我们在上面单元格中定义的 generator 对象。

您可以参考 此脚本 以了解有关如何将 LibriSpeech 数据转换为 tfrecords 的更多详细信息。

output_signature = (
    tf.TensorSpec(shape=(None),  dtype=tf.float32),
    tf.TensorSpec(shape=(None), dtype=tf.int32),
)

dataset = tf.data.Dataset.from_generator(inputs_generator, output_signature=output_signature)

BUFFER_SIZE = len(flac_files)
SEED = 42

dataset = dataset.shuffle(BUFFER_SIZE, seed=SEED)

我们将把数据集传递到多个批次中，因此让我们在以下单元格中准备批次。现在，批次中的所有序列都应填充到恒定长度。我们将为此目的使用 .padded_batch(...) 方法。

dataset = dataset.padded_batch(BATCH_SIZE, padded_shapes=(AUDIO_MAXLEN, LABEL_MAXLEN), padding_values=(0.0, 0))

加速器（如 GPU/TPU）非常快，并且在训练期间数据加载（和预处理）通常成为瓶颈，因为数据加载部分发生在 CPU 上。这会显着增加训练时间，尤其是在涉及大量在线预处理或数据从 GCS 存储桶在线流式传输时。为了解决这些问题，tf.data.Dataset 提供了 .prefetch(...) 方法。此方法有助于在模型对当前批次进行预测（在 GPU/TPU 上）的同时，并行（在 CPU 上）准备接下来的几个批次。

dataset = dataset.prefetch(tf.data.AUTOTUNE)

由于此笔记本是为演示目的而制作的，因此我们将采用前 num_train_batches 并仅对其执行训练。但是，鼓励您对整个数据集进行训练。类似地，我们只评估 num_val_batches。

num_train_batches = 10
num_val_batches = 4

train_dataset = dataset.take(num_train_batches)
val_dataset = dataset.skip(num_train_batches).take(num_val_batches)

模型训练

为了训练我们的模型，我们将在使用 .compile(...) 编译我们的模型后直接调用 .fit(...) 方法。

model.compile(optimizer, loss=loss_fn)

上面的单元格将设置我们的训练状态。现在，我们可以使用 .fit(...) 方法启动训练。

history = model.fit(train_dataset, validation_data=val_dataset, epochs=3)
history.history

让我们使用 .save(...) 方法保存我们的模型，以便以后能够执行推理。您还可以按照 TFHub 文档 将此 SavedModel 导出到 TFHub。

save_dir = "finetuned-wav2vec2"
model.save(save_dir, include_optimizer=False)

评估

现在，我们将计算验证数据集上的词错误率

词错误率 (WER) 是用于衡量自动语音识别系统性能的常用指标。WER 源自莱文斯坦距离，在词级别工作。词错误率可以计算为：WER = (S + D + I) / N = (S + D + I) / (S + D + C)，其中 S 是替换次数，D 是删除次数，I 是插入次数，C 是正确词语数，N 是参考中的词语数 (N=S+D+C)。此值表示错误预测的词语百分比。

您可以参考 这篇论文 以了解有关 WER 的更多信息。

我们将使用 HuggingFace 数据集 库中的 load_metric(...) 函数。让我们首先使用 pip 安装 datasets 库，然后定义 metric 对象。

!pip3 install -q datasets

from datasets import load_metric
metric = load_metric("wer")

@tf.function(jit_compile=True)
def eval_fwd(batch):
  logits = model(batch, training=False)
  return tf.argmax(logits, axis=-1)

现在是时候在验证数据上运行评估了。

from tqdm.auto import tqdm

for speech, labels in tqdm(val_dataset, total=num_val_batches):
    predictions  = eval_fwd(speech)
    predictions = [tokenizer.decode(pred) for pred in predictions.numpy().tolist()]
    references = [tokenizer.decode(label, group_tokens=False) for label in labels.numpy().tolist()]
    metric.add_batch(references=references, predictions=predictions)

我们使用 tokenizer.decode(...) 方法将我们的预测和标签解码回文本，并将它们添加到指标中，以便稍后进行 WER 计算。

现在，让我们在以下单元格中计算指标值

metric.compute()

推理

现在，我们对训练过程感到满意，并且已将模型保存在 save_dir 中，我们将看看如何使用此模型进行推理。

首先，我们将使用 tf.keras.models.load_model(...) 加载我们的模型。

finetuned_model = tf.keras.models.load_model(save_dir)

让我们下载一些语音样本以执行推理。您也可以用自己的语音样本替换以下样本。

wget https://github.com/vasudevgupta7/gsoc-wav2vec2/raw/main/data/SA2.wav

现在，我们将使用 soundfile.read(...) 读取语音样本，并将其填充到 AUDIO_MAXLEN 以满足模型签名。然后，我们将使用 Wav2Vec2Processor 实例对该语音样本进行归一化，并将其馈送到模型中。

import numpy as np

speech, _ = sf.read("SA2.wav")
speech = np.pad(speech, (0, AUDIO_MAXLEN - len(speech)))
speech = tf.expand_dims(processor(tf.constant(speech)), 0)

outputs = finetuned_model(speech)
outputs

让我们使用上面定义的 Wav2Vec2tokenizer 实例将数字解码回文本序列。

predictions = tf.argmax(outputs, axis=-1)
predictions = [tokenizer.decode(pred) for pred in predictions.numpy().tolist()]
predictions

此预测非常随机，因为模型从未在本笔记本中接受过大型数据的训练（因为本笔记本并非用于进行完整的训练）。如果您在完整的 LibriSpeech 数据集上训练此模型，您将获得良好的预测结果。

最后，我们已经完成了本笔记本。但这并不是学习 TensorFlow 用于语音相关任务的结束，此存储库 包含一些更棒的教程。如果您在本笔记本中遇到任何错误，请在 此处 创建一个问题。