Distil-Whisper

[论文] [模型] [Colab] [训练代码]

Distil-Whisper 是 Whisper 的蒸馏版英语语音识别模型，其速度提升了 6 倍，模型大小缩小了 49%，并且在分布外评估集上的词错误率（WER）仅比原模型高 1%：

对于大多数应用场景，我们推荐使用最新的 distil-large-v3 检查点，因为它是在所有 Whisper 库中兼容性最好、性能最强的蒸馏模型。唯一的例外是资源受限、内存极少的应用场景，例如设备端或移动应用，在这种情况下，distil-small.en 是一个很好的选择，因为它只有 1.66 亿参数，且在 WER 方面仅比 Whisper large-v3 高 4%。

[!注意]
Distil-Whisper 目前仅支持英语语音识别。对于多语言语音识别，我们建议使用由 OpenAI 发布的 Whisper Turbo 检查点，它采用了与 Distil-Whisper 相同的原理。有关详细信息，请参阅 Whisper Turbo 的 发布说明。

1. 使用方法

Distil-Whisper 自 Hugging Face 🤗 Transformers 4.35 版本起得到支持。要运行该模型，首先需要安装最新版本的 Transformers 库。在此示例中，我们还将安装 🤗 Datasets，以便从 Hugging Face Hub 加载一个玩具音频数据集：

pip install --upgrade pip
pip install --upgrade transformers accelerate datasets[audio]

短格式转录

短格式转录是指对长度小于30秒的音频片段进行转录的过程，这正是Whisper模型的最大接收范围。这意味着整个音频片段可以一次性处理完毕，而无需将其分割成多个部分。

首先，我们通过便捷的AutoModelForSpeechSeq2Seq和AutoProcessor类加载Distil-Whisper模型。

我们将模型以float16精度加载，并通过传递low_cpu_mem_usage=True来尽可能缩短加载时间。此外，我们还希望确保模型以safetensors格式加载，因此传递了use_safetensors=True：

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "distil-whisper/distil-large-v3"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

随后，可以将模型和处理器传递给pipeline。请注意，如果您希望对生成过程有更多控制，可以直接使用模型加处理器的API，如下所示。

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    torch_dtype=torch_dtype,
    device=device,
)

接下来，我们从LibriSpeech语料库中加载一个示例短格式音频：

from datasets import load_dataset

dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
sample = dataset[0]["audio"]

最后，我们可以调用管道来转录音频：

result = pipe(sample)
print(result["text"])

若要转录本地音频文件，只需在调用管道时传入音频文件的路径即可：

result = pipe("audio.mp3")
print(result["text"])

有关如何自定义自动语音识别管道的更多信息，请参阅ASR管道的文档。我们还提供了一个端到端的Google Colab，用于对比Whisper与Distil-Whisper的性能。

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
from datasets import Audio, load_dataset


device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "distil-whisper/distil-large-v3"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
dataset = dataset.cast_column("audio", Audio(processor.feature_extractor.sampling_rate))
sample = dataset[0]["audio"]

input_features = processor(
  sample["array"], sampling_rate=sample["sampling_rate"], return_tensors="pt"
).input_features

input_features = input_features.to(device, dtype=torch_dtype)

gen_kwargs = {
  "max_new_tokens": 128,
  "num_beams": 1,
  "return_timestamps": False,
}

pred_ids = model.generate(input_features, **gen_kwargs)
pred_text = processor.batch_decode(pred_ids, skip_special_tokens=True, decode_with_timestamps=gen_kwargs["return_timestamps"])

print(pred_text)

顺序长文本模式

最新的 distil-large-v3 检查点专门设计用于兼容 OpenAI 的顺序长文本转录算法。该算法采用滑动窗口对较长的音频文件（超过 30 秒）进行缓冲推理，与 分块长文本算法 相比，能够生成更准确的转录结果。

顺序长文本算法应在以下任一场景中使用：

1. 转录准确性是最重要的因素，而延迟相对不那么重要。
2. 您正在批量转录较长的音频文件，在这种情况下，顺序模式的延迟与分块模式相当，但字错误率（WER）可低至 0.5%。

如果您只转录单个较长的音频文件，并且延迟是最关键的因素，则应使用下文所述的分块算法。有关不同算法的详细说明，请参阅 Distil-Whisper 论文 的第 5 节。

首先，我们像之前一样加载模型和处理器：

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "distil-whisper/distil-large-v3"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

随后，可以将模型和处理器传递给 pipeline。请注意，如果您希望对生成过程有更多控制，可以直接使用 model.generate(...) API，如下所示。

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    torch_dtype=torch_dtype,
    device=device,
)

接下来，我们加载一个长文本音频样本。这里我们使用 LibriSpeech 语料库中的拼接示例：

from datasets import load_dataset

dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation")
sample = dataset[0]["audio"]

最后，我们可以调用管道来转录音频：

result = pipe(sample)
print(result["text"])

要转录音频文件，只需在调用管道时传入音频文件的路径即可：

result = pipe("audio.mp3")
print(result["text"])

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
from datasets import Audio, load_dataset


device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "distil-whisper/distil-large-v3"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
dataset = dataset.cast_column("audio", Audio(processor.feature_extractor.sampling_rate))
sample = dataset[0]["audio"]

inputs = processor(
    sample["array"],
    sampling_rate=sample["sampling_rate"],
    return_tensors="pt",
    truncation=False,
    padding="longest",
    return_attention_mask=True,
)
inputs = inputs.to(device, dtype=torch_dtype)

gen_kwargs = {
    "max_new_tokens": 448,
    "num_beams": 1,
    "condition_on_prev_tokens": False,
    "compression_ratio_threshold": 1.35,  # zlib 压缩比阈值（以 token 空间为单位）
    "temperature": (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
    "logprob_threshold": -1.0,
    "no_speech_threshold": 0.6,
    "return_timestamps": True,
}

pred_ids = model.generate(**inputs, **gen_kwargs)
pred_text = processor.batch_decode(pred_ids, skip_special_tokens=True, decode_with_timestamps=False)

print(pred_text)

分块长文本模式

distil-large-v3 仍然兼容 Transformers 的分块长文本算法。当需要转录单个大型音频文件且要求尽可能快的推理速度时，应使用此算法。在这种情况下，分块算法的速度可比 OpenAI 的顺序长文本实现快高达 9 倍（参见 Distil-Whisper 论文 中的表 7）。

我们可以像之前一样加载模型和处理器：

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "distil-whisper/distil-large-v3"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

要启用分块功能，需将 chunk_length_s 参数传递给 pipeline。对于 distil-large-v3，最佳的分块长度为 25 秒。若要激活批处理，则需传入 batch_size 参数：

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    chunk_length_s=25,
    batch_size=16,
    torch_dtype=torch_dtype,
    device=device,
)

max_new_tokens 参数控制每个分块生成的最大标记数。在典型的语音场景中，每秒通常不会超过 3 个词。因此，对于 30 秒的输入，最多约有 90 个词（约 128 个标记）。我们将每个分块的最大生成标记数设置为 128，以截断可能在片段末尾出现的幻觉内容。这些标记会在分块边界处通过长文本分块转录算法被移除，因此直接在生成过程中进行截断更为高效，可以避免解码器中多余的生成步骤。

现在，让我们加载一个长文本音频样本。这里我们使用 LibriSpeech 语料库中的拼接示例：

from datasets import load_dataset

dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation")
sample = dataset[0]["audio"]

最后，我们可以调用管道来转录音频：

result = pipe(sample)
print(result["text"])

有关如何自定义自动语音识别管道的更多信息，请参阅 ASR 管道的 文档。

推测解码

Distil-Whisper 可以作为 Whisper 的辅助模型，用于 推测解码。推测解码在数学上能够确保与 Whisper 完全相同的输出，同时速度提升至两倍。这使其成为现有 Whisper 管道的理想替代品，因为可以保证输出完全一致。

进行推测解码时，我们需要同时加载教师模型：openai/whisper-large-v3，以及辅助模型（即学生模型）：distil-whisper/distil-large-v3。

首先，我们加载教师模型及其处理器。加载方式与前面示例中加载 Distil-Whisper 模型的方式大致相同：

from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor
import torch

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "openai/whisper-large-v3"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

接下来，我们加载辅助模型。由于 Distil-Whisper 与教师模型共享完全相同的编码器，因此我们只需加载一个两层的解码器，将其视为“仅解码器”模型：

from transformers import AutoModelForCausalLM
assistant_model_id = "distil-whisper/distil-large-v2"

assistant_model = AutoModelForCausalLM.from_pretrained(
    assistant_model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
assistant_model.to(device)

辅助模型与教师模型共用同一处理器，因此无需单独加载学生处理器。

现在，我们可以将辅助模型传递给管道，以用于推测解码。我们将它作为 generate_kwarg 参数，键名为 "assistant_model"，从而启用推测解码功能：

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    generate_kwargs={"assistant_model": assistant_model},
    torch_dtype=torch_dtype,
    device=device,
)

与之前一样，我们可以将任意样本传递给管道进行转录：

from datasets import load_dataset

dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
sample = dataset[0]["audio"]

result = pipe(sample)
print(result["text"])

注意： 推测解码的平均速度应比仅使用 Whisper large-v2 快两倍，而显存占用仅增加约 8%，同时在数学上确保结果完全一致。这使得它成为现有语音识别管道中替换 Whisper large-v2 的理想选择。

有关推测解码的更多详细信息，请参阅以下资源：

• Sanchit Gandhi 撰写的博客文章 用于 2 倍更快 Whisper 推理的推测解码
• Joao Gante 撰写的博客文章 辅助生成：迈向低延迟文本生成的新方向
• Leviathan 等人撰写的论文 通过推测解码实现 Transformers 的快速推理

其他速度与内存优化

我们可以在下文中介绍一些额外的速度和内存优化方法，适用于 Distil-Whisper。

Flash Attention

如果您的 GPU 支持，我们建议使用 Flash Attention 2。为此，您需要先安装 Flash Attention：

pip install flash-attn --no-build-isolation

然后，您可以将 use_flash_attention_2=True 传递给 from_pretrained 方法，以启用 Flash Attention 2：

- model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
+ model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True, use_flash_attention_2=True)

Torch Scale-Product-Attention (SDPA)

如果您的 GPU 不支持 Flash Attention，我们建议使用 BetterTransformers。为此，您需要先安装 optimum：

pip install --upgrade optimum

然后，在使用模型之前将其转换为 "BetterTransformer" 模型：

model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
+ model = model.to_bettertransformer()

导出到其他库

Distil-Whisper 在以下库中支持原始的“顺序式”长文本转录算法。点击表格中的链接即可查看各库的相关代码片段：

随着“分块式”长文本转录算法集成到各个库中，相关更新将在此处发布。

有关 🤗 Transformers 的代码示例，请参阅“短文本转录”和“长文本转录”两节。

2. 为什么使用 Distil-Whisper？ ⁉️

Distil-Whisper 被设计为英语语音识别领域中 Whisper 的直接替代品。以下是切换到 Distil-Whisper 的 5 个理由：

1. 更快的推理速度： 推理速度是 Whisper 的 6 倍，同时在分布外音频上的 WER 误差率仅比 Whisper 高 1%：

2. 抗噪能力强： 在低信噪比条件下表现出色的 WER 性能：

3. 减少幻觉现象： 经量化分析，Distil-Whisper 的重复 5-gram 单词数量（5-Dup.）比 Whisper 少 1.3 倍，插入错误率（IER）也降低了 2.1%：

4. 专为推测解码设计： Distil-Whisper 可用作 Whisper 的辅助模型，使推理速度提升至原来的 2 倍，同时在数学上保证输出结果与 Whisper 相同。
5. 宽松的许可证： Distil-Whisper 采用 MIT 许可证，因此可用于商业用途。

3. 方法 ✍️

为了蒸馏 Whisper，我们将整个编码器模块复制并冻结，只保留两个解码器层，并从 Whisper 的第一个和最后一个解码器层初始化它们。其余的解码器层则被舍弃：

Distil-Whisper 是基于知识蒸馏的目标进行训练的。具体来说，它通过最小化蒸馏模型与 Whisper 模型之间的 KL 散度，以及伪标签音频数据上的交叉熵损失来进行训练。

我们使用总计 22,000 小时的伪标签音频数据对 Distil-Whisper 进行训练，这些数据涵盖了 10 个领域，涉及超过 18,000 名说话者：

如此多样化的音频数据集对于确保 Distil-Whisper 在不同数据集和领域中的鲁棒性至关重要。

此外，我们还使用 WER 过滤器来剔除 Whisper 错误转录或产生幻觉的伪标签，这大大提升了下游蒸馏模型的 WER 性能。

有关蒸馏设置和评估结果的详细信息，请参阅 Distil-Whisper 论文。

4. 训练代码

用于复现 Distil-Whisper 的训练代码可在 training 目录中找到。该代码经过调整，足以支持多语言语音识别的 Whisper 蒸馏工作，方便社区中的任何人根据自己的需求蒸馏任意语言的 Whisper。

5. 致谢

• OpenAI 提供的 Whisper 模型及原始代码库
• Hugging Face 🤗 Transformers 在模型集成方面的支持
• Google 的 TPU 研究云 (TRC) 计划提供的 Cloud TPU v4 资源

6. 引用

如果您使用本模型，请考虑引用 Distil-Whisper 论文：

@misc{gandhi2023distilwhisper,
      title={Distil-Whisper: Robust Knowledge Distillation via Large-Scale Pseudo Labelling}, 
      author={Sanchit Gandhi and Patrick von Platen and Alexander M. Rush},
      year={2023},
      eprint={2311.00430},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

同时请引用 Whisper 论文：

@misc{radford2022robust,
      title={Robust Speech Recognition via Large-Scale Weak Supervision}, 
      author={Alec Radford and Jong Wook Kim and Tao Xu and Greg Brockman and Christine McLeavey and Ilya Sutskever},
      year={2022},
      eprint={2212.04356},
      archivePrefix={arXiv},
      primaryClass={eess.AS}
}

Distil-Whisper 快速上手指南

Distil-Whisper 是 Whisper 模型的蒸馏版本，专为英语语音识别设计。相比原版，它速度快 6 倍，体积小 49%，且在分布外测试集上的词错误率（WER）差异控制在 1% 以内。

注意：本模型仅支持英语。如需多语言支持，建议使用 OpenAI 发布的 Whisper Turbo 模型。

1. 环境准备

• 系统要求：支持 Linux、macOS 或 Windows。
• 硬件建议：推荐使用 NVIDIA GPU 以获得最佳推理速度（需安装 CUDA）。若无 GPU，可在 CPU 上运行但速度较慢。
• Python 版本：建议 Python 3.8+。
• 核心依赖：
  • transformers >= 4.35
  • torch (PyTorch)
  • accelerate
  • datasets (用于加载测试音频)

2. 安装步骤

使用 pip 安装最新版本的依赖库。国内用户建议使用清华源或阿里源加速下载。

# 升级 pip
pip install --upgrade pip

# 安装核心依赖 (推荐使用国内镜像源)
pip install --upgrade transformers accelerate datasets[audio] -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 基本使用

以下示例演示如何加载 distil-large-v3 模型并对短音频（<30 秒）进行转录。该模型会自动检测是否有 GPU 并启用半精度 (float16) 加速。

最简单的使用示例

import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from datasets import load_dataset

# 1. 设置设备和数据类型
device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

# 2. 加载模型和处理器
# 推荐使用 distil-large-v3，它在性能和兼容性之间取得了最佳平衡
model_id = "distil-whisper/distil-large-v3"

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, 
    torch_dtype=torch_dtype, 
    low_cpu_mem_usage=True, 
    use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

# 3. 创建语音识别管道
pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    torch_dtype=torch_dtype,
    device=device,
)

# 4. 执行转录

# 方式 A: 转录在线测试数据集中的一个样本
dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
sample = dataset[0]["audio"]
result = pipe(sample)
print("在线样本转录结果:", result["text"])

# 方式 B: 转录本地音频文件 (取消注释并替换文件名即可使用)
# result = pipe("path/to/your/audio.mp3")
# print("本地文件转录结果:", result["text"])

长音频处理提示

如果需要转录超过 30 秒的长音频，distil-large-v3 支持顺序长形式（Sequential Long-Form）算法。只需将音频文件路径传入上述 pipe 即可，管道会自动处理滑动窗口推理，无需手动切片。