Whisper 模型微调全指南

OpenAI Whisper 是一个基于 Transformer 的语音转文字模型，有着优秀的讲稿识别能力。默认模型对于大多数语言来说足够精确，但仍然可以通过微调来取得更好的性能。

本文将介绍微调 Whisper 的全过程：如何准备数据集、微调过程本身、以及如何高效地部署微调后的模型。

准备相关工具

首先，你需要安装一些工具

Pip 依赖库

Bash

pip install transformers datasets huggingface-hub accelerate evaluate tensorboard

登录到 HuggingFace

在安装 pip 依赖库后，我们需要登入 HuggingFace。

Bash

huggingface-cli login

登录需要 token. 在上面这条指令中，huggingface-cli 会提供一个网址，告诉你如何获取你的 API token。如果需要将模型上传到 HuggingFace Hub，你需要一个拥有 write 权限的 token。

准备数据集

Whisper 是一个监督学习的模型。因此在数据集中，你需要提供音频文件以及音频对应的文字。

最简单的数据集准备方法是使用 HuggingFace AudioFolder. 建立文件夹，并将你的文件如下摆放：

Plaintext

folder/train/metadata.jsonl
folder/train/first.mp3
folder/train/second.mp3
folder/train/third.mp3

注意不是所有文件都支持。例如，m4a 文件就无法使用。为了文件整理方便，文件可以放到子文件夹中。

metadata.jsonl 是一个 JSON Lines 格式的文件，其格式如下。

JSONL

{"file_name": "first.mp3", "transcription": "First Audio Transcription"}
{"file_name": "second.mp3", "transcription": "Second Audio Transcription"}
{"file_name": "third.mp3", "transcription": "Third Audio Transcription"}

JSONL 的意思是 JSON Lines: 每一行是一个 JSON 对象，而整个文件可以被看作是一个数组的 JSON 对象。
在每行中，file_name 的名字必须是 file_name。它提供的是音频文件的相对路径，相对这一个 metadata.jsonl文件。
其他键值（比如transcription）可以任意起名。最后这个文件将会被转成 Arrow 格式的表格（类似 pandas 的 Dataset），而每一个键值对应的表格中的一列。
你可以加入任意多的其他键值，可以指明说话人、语言、来源等信息。

数据集准备完成后，使用如下命令将数据集上传至 HuggingFace Hub.

Python

from datasets import load_dataset
audio_dataset = load_dataset("audiofolder", data_dir=".")
audio_dataset.push_to_hub("YOUR_HF_NAME/HF_DATASET_REPO") # Replace this with your Huggingface Repository

这将会读取音频文件，将整个数据集转换为 Parquet 格式，自动生成包含数据集信息的 README.md 文件，并上传到 HuggingFace Hub.

微调

微调基于 HuggingFace 版本的 OpenAI Whisper 模型。关于微调的详细过程可以在这里找到 https://huggingface.co/blog/fine-tune-whisper.

此处我们提供包含了少许修改的源代码。

Python

# NOTE: 注意：在此处填入finetune 的基座模型。
base_model = "openai/whisper-large-v3"

# NOTE: 此处不要修改。除非你想训练 translate 模式，且你的数据集包含原音频的英文翻译。
task = "transcribe"

from datasets import load_dataset, DatasetDict

# ========== Load Dataset ==========
tl_dataset = DatasetDict()
tl_dataset["train"] = load_dataset("YOUR_HF_NAME/HF_DATASET_REPO", split="train")
# NOTE: 如果你的数据集包含 test 分区，将下一行取消注释
# tl_dataset["test"] = load_dataset("metricv/tl-whisper", "hi", split="test")

# ========== Load Whisper Preprocessor ==========

from transformers import WhisperFeatureExtractor
from transformers import WhisperTokenizer
from transformers import WhisperProcessor

feature_extractor = WhisperFeatureExtractor.from_pretrained(base_model)
tokenizer = WhisperTokenizer.from_pretrained(base_model, task=task)
processor = WhisperProcessor.from_pretrained(base_model, task=task)

# ========== Process Dataset ==========

from datasets import Audio

tl_dataset = tl_dataset.cast_column("audio", Audio(sampling_rate=16000))

def prepare_dataset(batch):
    # load and resample audio data from 48 to 16kHz
    audio = batch["audio"]
    # compute log-Mel input features from input audio array
    batch["input_features"] = feature_extractor(audio["array"], sampling_rate=audio["sampling_rate"]).input_features[0]
    # encode target text to label ids
    # NOTE: 此处的键值 "transcription" 指的就是你在创建数据集的过程中，包含音频文件对应文字的键值。如果你是用的键名不是 transcription，在此处修改。
    batch["labels"] = tokenizer(batch["transcription"]).input_ids
    return batch

tl_dataset = tl_dataset.map(prepare_dataset, remove_columns=tl_dataset.column_names["train"], num_proc=8)

# ========== Load Whisper Model ==========

from transformers import WhisperForConditionalGeneration
model = WhisperForConditionalGeneration.from_pretrained(base_model)
model.generation_config.task = task
model.generation_config.forced_decoder_ids = None

# ========== Fine-tune model ==========

import torch
from dataclasses import dataclass
from typing import Any, Dict, List, Union

@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any
    decoder_start_token_id: int

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        # split inputs and labels since they have to be of different lengths and need different padding methods
        # first treat the audio inputs by simply returning torch tensors
        input_features = [{"input_features": feature["input_features"]} for feature in features]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")

        # get the tokenized label sequences
        label_features = [{"input_ids": feature["labels"]} for feature in features]
        # pad the labels to max length
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")

        # replace padding with -100 to ignore loss correctly
        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)

        # if bos token is appended in previous tokenization step,
        # cut bos token here as it's append later anyways
        if (labels[:, 0] == self.decoder_start_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels

        return batch
    
data_collator = DataCollatorSpeechSeq2SeqWithPadding(
    processor=processor,
    decoder_start_token_id=model.config.decoder_start_token_id,
)

import evaluate

metric = evaluate.load("wer")

def compute_metrics(pred):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    # replace -100 with the pad_token_id
    label_ids[label_ids == -100] = tokenizer.pad_token_id

    # we do not want to group tokens when computing the metrics
    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    wer = 100 * metric.compute(predictions=pred_str, references=label_str)

    return {"wer": wer}

from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments

training_args = Seq2SeqTrainingArguments(
    output_dir="./whisper-large-v3-ft-train",  # change to a repo name of your choice
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,  # increase by 2x for every 2x decrease in batch size
    learning_rate=1e-5,
    num_train_epochs=2.0,
    # warmup_steps=500,
    # max_steps=4000,
    gradient_checkpointing=True,
    fp16=True,
    do_eval=False,
    # eval_strategy="steps",    # NOTE: 如果你的数据集包含 test 分区，可取消注释此行
    per_device_eval_batch_size=8,
    predict_with_generate=True,
    generation_max_length=225,
    # save_steps=1000,
    # eval_steps=1000,
    logging_steps=5,
    report_to=["tensorboard"],
    load_best_model_at_end=False,
    metric_for_best_model="wer",
    greater_is_better=False,
    push_to_hub=False,
)

trainer = Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=tl_dataset["train"],
    # eval_dataset=tl_dataset["test"], # NOTE: 如果你的数据集包含 test 分区，可取消注释此行
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    tokenizer=processor.feature_extractor,
)

processor.save_pretrained(training_args.output_dir)

trainer.train()

# ========== Save model ==========

trainer.save_model(output_dir="./whisper-large-v3-finetuned")
torch.save(model.state_dict(), f"{training_args.output_dir}/pytorch_model.bin")

# ========== Push model to HF hub ==========
# 如果你不想上传模型，注释掉以下行。
trainer.push_to_hub("YOUR_HF_NAME/HF_MODEL_REPO") # 修改为你的 HuggingFace 仓库名

注意这段代码使用了以下几个文件夹

训练过程文件夹: ./whisper-large-v3-ft-train
模型输出文件夹: ./whisper-large-v3-finetuned

准备工作

在部署微调后的模型前，你还需要设置一些东西

基座模型的 tokenizer.json 可能没有被复制过来。你需要手动复制一下。
- 在 HuggingFace Hub上找到基座模型 (such as openai/whisper-large-v3), 下载它的 tokenizer.json，并放到“模型输出文件夹”下。
- 如果你使用了 push_to_hub() 来上传模型，但是上传后的模型没有 tokenizer.json，你可以使用 HuggingFace 的网页界面手动上传。
- 如果“模型输出文件夹”中没有 tokenizer_config.json，将“训练过程文件夹”的对应文件复制过来。
- 如果“模型输出文件夹”中没有 preprocessor_config.json，将“训练过程文件夹”的对应文件复制过来

部署

微调后的模型可以用多种方式部署。

使用 HuggingFace 运行库

微调后的模型和原版 HuggingFace 模型一样，可以使用 from_pretrained() 来部署。此方法略快于原版 openai-whisper 包，但会占用更多 RAM。

Python

from transformers import WhisperForConditionalGeneration, WhisperProcessor

processor = WhisperProcessor.from_pretrained("YOUR_HF_NAME/HF_MODEL_REPO")
# 如果你没有上传模型，使用 from_pretrained("模型输出文件夹") 加载本地模型。

使用原版 PyPI 的 `openai-whisper` 包

微调后的模型可被转换为兼容原版 openai-whisper 包的格式。

在微调的结尾，我们把 PyTorch 格式的模型保存在了 "训练过程文件夹"/pytorch_model.bin. 但是，这个模型中的层和原版模型的命名方式不一样。重命名即可解决该问题。

使用以下代码转换，来自 https://github.com/openai/whisper/discussions/830

Python

# NOTE: Change this to the base model you fine-tuned from.
BASE_MODEL = "large-v3"

#!/bin/env python3
import whisper
import re
import torch

def hf_to_whisper_states(text):
    text = re.sub('.layers.', '.blocks.', text)
    text = re.sub('.self_attn.', '.attn.', text)
    text = re.sub('.q_proj.', '.query.', text)
    text = re.sub('.k_proj.', '.key.', text)
    text = re.sub('.v_proj.', '.value.', text)
    text = re.sub('.out_proj.', '.out.', text)
    text = re.sub('.fc1.', '.mlp.0.', text)
    text = re.sub('.fc2.', '.mlp.2.', text)
    text = re.sub('.fc3.', '.mlp.3.', text)
    text = re.sub('.fc3.', '.mlp.3.', text)
    text = re.sub('.encoder_attn.', '.cross_attn.', text)
    text = re.sub('.cross_attn.ln.', '.cross_attn_ln.', text)
    text = re.sub('.embed_positions.weight', '.positional_embedding', text)
    text = re.sub('.embed_tokens.', '.token_embedding.', text)
    text = re.sub('model.', '', text)
    text = re.sub('attn.layer_norm.', 'attn_ln.', text)
    text = re.sub('.final_layer_norm.', '.mlp_ln.', text)
    text = re.sub('encoder.layer_norm.', 'encoder.ln_post.', text)
    text = re.sub('decoder.layer_norm.', 'decoder.ln.', text)
    text = re.sub('proj_out.weight', 'decoder.token_embedding.weight', text)
    return text

# Load HF Model
# NOTE: Change the following line to point to "Training Data Directory"/pytorch_model.bin
hf_state_dict = torch.load("Training Data Directory/pytorch_model.bin", map_location=torch.device('cpu'))

# Rename layers
for key in list(hf_state_dict.keys())[:]:
    new_key = hf_to_whisper_states(key)
    hf_state_dict[new_key] = hf_state_dict.pop(key)

model = whisper.load_model(BASE_MODEL)
dims = model.dims

# Save it
# NOTE: This will save file to whisper-model.bin. Change the path as you wish.
torch.save({
    "dims": model.dims.__dict__,
    "model_state_dict": hf_state_dict
}, "whisper-model.bin")

然后，你就可以使用原版的 whisper.load("whisper-model.bin") 来加载模型了。

Faster-Whisper (CTranslate 2)

最有效率的部署方式是使用 faster-whisper 运行库，但我们需要再转换一次格式。

首先，安装faster-whisper的转换器

Bash

git clone --depth=1 https://github.com/SYSTRAN/faster-whisper
cd faster-whisper
pip install -e .[convert] # In zsh, quote ".[convert]"

然后使用以下命令进行转换

Bash

ct2-transformers-converter \
    --model YOUR_HF_NAME/HF_MODEL_REPO \
    --output_dir whisper-largve-v3-ft-ct2-f16 \
    --copy_files tokenizer.json preprocessor_config.json \
    --quantization float16

CTranslate2 模型会保存到一个文件夹，而并不是单一文件。将 whisper-largve-v3-ft-ct2-f16 改为你的目标文件夹。
Quantization 不是必要的。我们训练时使用的就是f16，所以此处的 quantization 其实并没做任何量化。

然后，微调后的模型就可以像任何其他模型一样，被 faster-whisper 加载。

Bash

from faster_whisper import WhisperModel

model = WhisperModel("/path/to/model/directory", device="cuda", compute_type="float16")

微调快乐～

Whisper 模型微调全指南

准备相关工具

Pip 依赖库

登录到 HuggingFace

准备数据集

微调

准备工作

部署

使用 HuggingFace 运行库

使用原版 PyPI 的 `openai-whisper` 包

Faster-Whisper (CTranslate 2)

语法自然的词典生成方式 – 以邦布语为例

Comments 1 条评论

博主 Agree

取消回复

准备相关工具

Pip 依赖库

登录到 HuggingFace

准备数据集

微调

准备工作

部署

使用 HuggingFace 运行库

使用原版 PyPI 的 openai-whisper 包

Faster-Whisper (CTranslate 2)

语法自然的词典生成方式 – 以邦布语为例

Comments 1 条评论

取消回复

使用原版 PyPI 的 `openai-whisper` 包