微調搭配 LM Head 的 Wav2Vec2

在 TensorFlow.org 上檢視

在 Google Colab 中執行

在 GitHub 上檢視

下載筆記本

查看 TF Hub 模型

在本筆記本中，我們將從 TFHub 載入預先訓練的 wav2vec2 模型，並透過在預先訓練模型之上附加語言模型化 Head (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 只需 1 行程式碼即可完成這兩項操作。

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 的 Functional API 包裝 pretrained_layer 和密集層 (LM head)。

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 源自 Levenshtein 距離，在詞層級運作。然後可以將詞錯誤率計算為：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 語音相關任務學習的結束，此 儲存庫 包含更多精彩的教學課程。如果您在此筆記本中遇到任何錯誤，請在此處建立問題 here。