Whisper-large-v3模型剪枝教程：减小模型大小保持精度-平芜编程栈

Whisper-large-v3模型剪枝教程：减小模型大小保持精度

你是不是也遇到过这种情况：想把一个强大的语音识别模型，比如Whisper-large-v3，塞进你的边缘设备里，结果发现它太大了，根本装不下？或者就算勉强装进去了，推理速度慢得像蜗牛，完全没法用。

我之前在做一个智能会议记录设备的时候，就卡在了这一步。Whisper-large-v3的识别效果确实好，但它的模型文件动辄几个GB，对内存和算力要求都很高。直接部署到树莓派或者Jetson这类设备上，要么内存爆掉，要么识别一句话要等十几秒，用户体验极差。

后来我花了不少时间研究模型剪枝，终于找到了一套比较靠谱的方法，能在显著减小模型体积的同时，尽量保住它的识别精度。今天我就把这套方法整理出来，手把手带你走一遍。你不用有太深的机器学习背景，跟着步骤做就行，目标是让你也能在资源有限的设备上跑起一个“瘦身”版的Whisper。

1. 准备工作：理解剪枝和搭建环境

在动手之前，我们先花几分钟搞清楚两件事：我们要剪的是什么，以及需要准备哪些工具。

1.1 模型剪枝到底在剪什么？

你可以把Whisper-large-v3这样的神经网络想象成一张非常复杂、密密麻麻的渔网。这张网由无数个“神经元”（也就是参数）连接而成，负责把输入的音频信号转换成文字。

模型剪枝，简单说，就是在这张网上找到那些不怎么重要的连接，然后把它剪掉。比如，有些连接权重值非常小，接近于零，说明它对最终结果的贡献微乎其微，剪掉它影响不大。剪掉这些冗余部分后，模型就变小了，计算量也少了，自然就跑得更快、更省资源。

我们的目标很明确：在精度下降可接受的范围内，让模型变得尽可能小、尽可能快。

1.2 搭建你的实验环境

工欲善其事，必先利其器。我们需要一个Python环境，并安装几个关键的库。我强烈建议使用Conda来管理环境，能避免很多依赖冲突的坑。

打开你的终端（Linux/Mac）或Anaconda Prompt（Windows），执行以下命令：

# 创建一个新的Python环境，命名为whisper_prune conda create -n whisper_prune python=3.10 -y conda activate whisper_prune # 安装PyTorch（请根据你的CUDA版本选择，如果没有GPU，去掉后面的CUDA部分） # 以下命令适用于CUDA 11.8，你可以在PyTorch官网找到适合你版本的命令 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 安装我们需要的核心库：Transformers（用于加载Whisper），以及剪枝工具库 pip install transformers datasets accelerate pip install torch-pruning # 这是一个非常好用的模型剪枝工具库

安装完成后，我们可以写个简单的测试脚本，确保能正常加载Whisper-large-v3模型。

# test_load.py - 测试模型加载 import torch from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor print("测试模型加载...") device = "cuda" if torch.cuda.is_available() else "cpu" print(f"使用设备: {device}") # 加载模型和处理器 model_id = "openai/whisper-large-v3" model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id) processor = AutoProcessor.from_pretrained(model_id) # 将模型移到指定设备 model.to(device) model.eval() # 设置为评估模式 print("模型加载成功！") print(f"模型参数量: {sum(p.numel() for p in model.parameters()):,}")

运行这个脚本，如果看到输出了巨大的参数量（大约15亿），恭喜你，环境搭建成功，我们有了一个“原版”的Whisper-large-v3作为起点。

2. 第一步：评估原始模型的性能和大小

在动剪刀之前，我们得先知道这把“尺子”原来有多长、多重。我们需要两样东西：一个测试数据集，和一套评估指标。

2.1 准备一个小的测试集

我们不需要用整个LibriSpeech那么大的数据集，那样太耗时。我们可以用Hugging Face Datasets库里的一个小样本集，或者自己准备几段音频。这里我用一个内置的小数据集做演示：

# evaluate_baseline.py - 评估原始模型 import torch from transformers import pipeline from datasets import load_dataset import numpy as np # 1. 加载模型和管道 device = "cuda" if torch.cuda.is_available() else "cpu" model_id = "openai/whisper-large-v3" pipe = pipeline( "automatic-speech-recognition", model=model_id, device=device, chunk_length_s=30 # 处理长音频时切片 ) # 2. 加载一个小型测试数据集（这里用Distil-Whisper提供的一个干净样本） print("加载测试数据...") dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation") # 我们只取前5个样本来快速评估 test_samples = dataset.select(range(5)) # 3. 定义一个简单的评估函数（计算词错误率WER） def compute_wer(reference, hypothesis): """一个简化的词错误率计算（实际应用建议使用专门的库如jiwer）""" ref_words = reference.lower().split() hyp_words = hypothesis.lower().split() # 这里简化处理，实际WER计算涉及动态规划 # 为了教程清晰，我们先用一个简单的差异比例 if len(ref_words) == 0: return 0.0 if len(hyp_words) == 0 else 1.0 # 简单计算匹配上的单词比例 matches = sum(1 for r, h in zip(ref_words, hyp_words) if r == h) return 1.0 - (matches / max(len(ref_words), len(hyp_words))) # 4. 运行评估 print("开始评估原始模型...") total_wer = 0 for i, sample in enumerate(test_samples): audio = sample["audio"] reference_text = sample["text"] # 使用模型进行识别 result = pipe(audio) hypothesis_text = result["text"] # 计算WER wer = compute_wer(reference_text, hypothesis_text) total_wer += wer print(f"样本 {i+1}:") print(f" 参考: {reference_text[:50]}...") print(f" 识别: {hypothesis_text[:50]}...") print(f" WER: {wer:.4f}") avg_wer = total_wer / len(test_samples) print(f"\n原始模型平均WER: {avg_wer:.4f}") # 5. 查看模型大小 import os from transformers import AutoModelForSpeechSeq2Seq model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id) model.save_pretrained("./original_whisper") original_size = sum(os.path.getsize(f) for f in os.listdir("./original_whisper") if f.endswith('.bin') or f.endswith('.safetensors')) print(f"原始模型文件大小: {original_size / 1024**3:.2f} GB")

运行这个脚本，你会得到原始模型在测试集上的平均词错误率（WER）和模型文件大小。记下这两个数字，这是我们剪枝后要对比的基准。

3. 第二步：选择并实施剪枝策略

剪枝有很多种方法，我们今天用两种比较常见且有效的方法结合：结构化剪枝和基于重要性的剪枝。

3.1 结构化剪枝：剪掉整个“神经元”

结构化剪枝不是剪掉单个连接，而是剪掉整个通道（channel）、注意力头（attention head）或者甚至整层（layer）。这就像不是剪掉渔网的一根线，而是直接抽掉一整排网眼。这样做的好处是，剪枝后的模型结构仍然是规则的，更容易被硬件加速。

对于Whisper这样的Transformer模型，注意力头（Attention Head）和FFN层（前馈网络）的中间维度是常见的剪枝目标。

# prune_structured.py - 结构化剪枝示例 import torch import torch.nn as nn from transformers import AutoModelForSpeechSeq2Seq import torch_pruning as tp # 导入剪枝工具库 # 1. 加载模型 model_id = "openai/whisper-large-v3" model = AutoModelForSpeechSeq2Seq.from_pretrained(model_id) model.eval() # 2. 定义要剪枝的层类型 # 我们主要针对注意力机制和前馈网络进行剪枝 def get_layers_to_prune(model): layers = [] # 遍历模型的编码器（encoder）和解码器（decoder） for module in model.modules(): if isinstance(module, nn.Linear): # 线性层是主要的剪枝目标 layers.append(module) return layers # 3. 构建剪枝依赖图（工具库会自动分析层之间的依赖关系） example_input = torch.randn(1, 80, 3000) # 模拟音频特征输入 DG = tp.DependencyGraph() DG.build_dependency(model, example_input=example_input) # 4. 选择剪枝策略：基于L1范数的重要性评分 # 权重绝对值越小的神经元，被认为越不重要 pruning_layers = get_layers_to_prune(model) pruning_idxs = {} # 记录每层要剪掉哪些索引 for layer in pruning_layers: if hasattr(layer, 'weight'): weight = layer.weight.data # 计算每个输出通道（神经元）的L1范数 importance = weight.abs().sum(dim=1) # 对行求和 # 假设我们想剪掉该层50%的通道 prune_ratio = 0.5 num_prune = int(len(importance) * prune_ratio) # 找到最不重要的通道索引 _, idxs = torch.topk(importance, k=num_prune, largest=False) pruning_idxs[layer] = idxs.tolist() # 5. 执行剪枝 pruning_plan = [] for layer, idxs in pruning_idxs.items(): pruning_plan.append(DG.get_pruning_plan(layer, tp.prune_linear_out_channel, idxs)) for plan in pruning_plan: plan.exec() # 执行剪枝计划 print("结构化剪枝完成！") print(f"剪枝后参数量: {sum(p.numel() for p in model.parameters()):,}") # 6. 保存剪枝后的模型 model.save_pretrained("./whisper_pruned_structured")

注意：上面的代码中，我们粗暴地剪掉了每个线性层50%的输出通道。在实际操作中，你需要更谨慎：

不要一次性剪太多：可以从10%-20%开始，评估精度损失。
不同层敏感度不同：靠近输入的层和靠近输出的层通常更关键，可以少剪或不剪。
需要微调（Fine-tuning）：剪枝后模型性能通常会下降，必须用数据重新训练（微调）一下，让模型适应新的结构。

3.2 基于重要性的非结构化剪枝：剪掉细小的连接

非结构化剪枝更精细，它针对的是单个权重参数。我们把网络中那些绝对值很小的权重（比如接近0的）直接置零。这就像在渔网上找到那些松垮无力的线头，把它剪断。这种剪枝能获得很高的稀疏度，但剪枝后的模型是不规则的稀疏矩阵，需要特殊的运行时库（如DeepSpeed、TensorRT）才能获得实际的加速。

# prune_unstructured.py - 非结构化剪枝（全局幅度剪枝） import torch import torch.nn as nn from transformers import AutoModelForSpeechSeq2Seq # 1. 加载模型（可以用剪枝过的，也可以用原始模型） model = AutoModelForSpeechSeq2Seq.from_pretrained("./whisper_pruned_structured") model.eval() # 2. 收集所有权重 all_weights = [] for name, param in model.named_parameters(): if 'weight' in name and param.dim() > 1: # 只剪枝多维权重，忽略bias等 all_weights.append(param.data.view(-1)) all_weights = torch.cat(all_weights) # 3. 确定全局阈值 # 例如，我们想剪掉全网络80%的最小权重 prune_ratio = 0.8 threshold = torch.quantile(all_weights.abs(), prune_ratio) print(f"全局剪枝阈值（绝对值）: {threshold.item()}") # 4. 应用剪枝掩码（将小于阈值的权重置零） total_params = 0 zero_params = 0 for name, param in model.named_parameters(): if 'weight' in name and param.dim() > 1: total_params += param.numel() # 创建掩码：绝对值大于阈值的为1，否则为0 mask = (param.data.abs() > threshold).float() zero_params += (mask == 0).sum().item() param.data.mul_(mask) # 应用掩码，小权重归零 sparsity = zero_params / total_params print(f"非结构化剪枝完成，稀疏度: {sparsity:.4f} ({zero_params}/{total_params} 个参数被置零)") # 5. 保存剪枝后的模型（注意：此时模型是稀疏的，保存的文件大小不会变） # 要获得实际的大小减少，需要将模型转换为稀疏格式或进行量化 model.save_pretrained("./whisper_pruned_final")

执行完这两步，我们得到了一个既被剪掉了整体结构（通道），内部又有很多零权重（稀疏）的模型。但它的精度很可能已经惨不忍睹了。接下来最关键的一步来了：微调。

4. 第三步：微调剪枝后的模型

剪枝相当于给模型做了一次大手术，术后必须进行康复训练（微调），让它重新学习如何利用剩下的“器官”正常工作。

4.1 准备微调数据

我们不需要海量数据，但需要一些有代表性的数据。可以使用原始Whisper训练数据的一小部分，或者你自己业务场景的音频数据。

# finetune.py - 微调剪枝后的模型 import torch from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, Seq2SeqTrainingArguments, Seq2SeqTrainer from datasets import load_dataset, Audio import numpy as np # 1. 加载剪枝后的模型和处理器 model_path = "./whisper_pruned_final" model = AutoModelForSpeechSeq2Seq.from_pretrained(model_path) processor = AutoProcessor.from_pretrained("openai/whisper-large-v3") # 2. 加载并预处理微调数据集（这里用LibriSpeech的一小部分） print("加载微调数据集...") dataset = load_dataset("librispeech_asr", "clean", split="train.100") # 只用100个样本演示 # 预处理函数：将音频转换为模型输入格式 def prepare_dataset(batch): audio = batch["audio"] # 重采样到16kHz（Whisper的标准输入） if audio["sampling_rate"] != 16000: # 这里需要音频处理库，如torchaudio或librosa # 为简化，假设数据已是16kHz pass batch["input_features"] = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").input_features[0] # 处理文本标签 batch["labels"] = processor.tokenizer(batch["text"]).input_ids return batch # 应用预处理 dataset = dataset.map(prepare_dataset, remove_columns=dataset.column_names) # 3. 分割训练集和验证集 split_dataset = dataset.train_test_split(test_size=0.2) train_dataset = split_dataset["train"] eval_dataset = split_dataset["test"] # 4. 定义数据整理器（Data Collator） from dataclasses import dataclass from typing import Any, Dict, List, Union @dataclass class DataCollatorSpeechSeq2SeqWithPadding: processor: Any def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]: input_features = [{"input_features": feature["input_features"]} for feature in features] label_features = [{"input_ids": feature["labels"]} for feature in features] # 填充输入特征 batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt") # 填充标签 labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt") # 将填充的标签token替换为-100，以便在损失计算时忽略 labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100) batch["labels"] = labels return batch data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor) # 5. 配置训练参数（关键步骤！） # 对于剪枝后的模型，学习率可以设小一点，训练轮数（epoch）也不用太多 training_args = Seq2SeqTrainingArguments( output_dir="./whisper_pruned_finetuned", per_device_train_batch_size=2, # 根据你的GPU内存调整 per_device_eval_batch_size=2, gradient_accumulation_steps=4, # 模拟更大的批次 learning_rate=1e-5, # 较小的学习率 warmup_steps=50, num_train_epochs=3, # 微调3轮通常足够 logging_dir="./logs", logging_steps=10, evaluation_strategy="epoch", save_strategy="epoch", load_best_model_at_end=True, metric_for_best_model="wer", # 根据WER选择最佳模型 greater_is_better=False, push_to_hub=False, # 设为True可上传到Hugging Face Hub ) # 6. 定义评估指标（词错误率WER） import evaluate metric = evaluate.load("wer") def compute_metrics(pred): pred_ids = pred.predictions label_ids = pred.label_ids # 将token id解码为文本 pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True) label_ids[label_ids == -100] = processor.tokenizer.pad_token_id label_str = processor.batch_decode(label_ids, skip_special_tokens=True) # 计算WER wer = 100 * metric.compute(predictions=pred_str, references=label_str) return {"wer": wer} # 7. 创建Trainer并开始微调 trainer = Seq2SeqTrainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, data_collator=data_collator, compute_metrics=compute_metrics, ) print("开始微调剪枝后的模型...") trainer.train() # 8. 保存最终模型 trainer.save_model("./whisper_pruned_finetuned_final") processor.save_pretrained("./whisper_pruned_finetuned_final") print("微调完成，模型已保存！")

微调过程可能需要一些时间，具体取决于你的数据量和GPU性能。耐心等待完成后，你就得到了一个“康复”后的精简版Whisper模型。

5. 第四步：评估剪枝效果与部署优化

最后，我们来验收一下成果，看看剪枝到底带来了多少好处，以及如何部署这个瘦身模型。

5.1 全面评估：精度、大小、速度

让我们写一个完整的评估脚本，对比原始模型和剪枝微调后模型的各项指标。

# final_evaluation.py - 最终评估 import torch from transformers import pipeline, AutoModelForSpeechSeq2Seq from datasets import load_dataset import time import os def evaluate_model(model_path, test_samples, device="cuda"): """评估指定模型的WER、推理速度和大小""" # 加载模型和管道 print(f"\n评估模型: {model_path}") pipe = pipeline( "automatic-speech-recognition", model=model_path, device=device, torch_dtype=torch.float16 if device=="cuda" else torch.float32 ) # 评估精度（WER） total_wer = 0 total_time = 0 for i, sample in enumerate(test_samples): audio = sample["audio"] reference_text = sample["text"] start_time = time.time() result = pipe(audio) end_time = time.time() hypothesis_text = result["text"] wer = compute_wer(reference_text, hypothesis_text) # 复用之前的函数 total_wer += wer total_time += (end_time - start_time) if i < 2: # 打印前两个样本的识别结果 print(f" 样本{i+1}识别: {hypothesis_text[:60]}...") avg_wer = total_wer / len(test_samples) avg_latency = total_time / len(test_samples) # 评估模型大小 model_size = 0 for file in os.listdir(model_path): if file.endswith('.bin') or file.endswith('.safetensors'): model_size += os.path.getsize(os.path.join(model_path, file)) return { "wer": avg_wer, "latency_seconds": avg_latency, "size_gb": model_size / 1024**3 } # 加载测试数据（用5个样本） dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation") test_samples = dataset.select(range(5)) # 评估原始模型 original_results = evaluate_model("openai/whisper-large-v3", test_samples) # 评估我们的剪枝微调模型 pruned_results = evaluate_model("./whisper_pruned_finetuned_final", test_samples) # 打印对比结果 print("\n" + "="*50) print("模型剪枝效果对比") print("="*50) print(f"{'指标':<20} {'原始模型':<15} {'剪枝模型':<15} {'变化':<10}") print(f"{'平均WER':<20} {original_results['wer']:<15.4f} {pruned_results['wer']:<15.4f} {pruned_results['wer']-original_results['wer']:+.4f}") print(f"{'推理延迟(秒)':<20} {original_results['latency_seconds']:<15.2f} {pruned_results['latency_seconds']:<15.2f} {((pruned_results['latency_seconds']/original_results['latency_seconds'])-1)*100:+.1f}%") print(f"{'模型大小(GB)':<20} {original_results['size_gb']:<15.2f} {pruned_results['size_gb']:<15.2f} {((pruned_results['size_gb']/original_results['size_gb'])-1)*100:+.1f}%") print("="*50)

运行这个脚本，你会看到一个清晰的对比表格。理想情况下，模型大小和推理延迟应该有显著下降（比如减少30%-50%），而WER的上升控制在可接受的范围内（比如上升不超过2-3个百分点）。

5.2 部署到边缘设备的建议

得到精简模型后，如何把它部署到树莓派、Jetson Nano或手机这类边缘设备上呢？这里有几个实用的建议：

进一步量化（Quantization）：剪枝之后，还可以对模型进行量化，将FP32的权重转换为INT8甚至INT4，这能再次大幅减少模型体积和提升推理速度。PyTorch提供了torch.quantization模块，可以尝试动态量化或静态量化。
```
# 简单的动态量化示例 import torch.quantization quantized_model = torch.quantization.quantize_dynamic( model, {torch.nn.Linear}, dtype=torch.qint8 )
```
转换为高效运行时格式：
- ONNX + ONNX Runtime：将PyTorch模型导出为ONNX格式，然后使用ONNX Runtime进行推理，它在CPU上通常有更好的优化。
- TensorRT：如果你有NVIDIA的Jetson设备，将模型转换为TensorRT引擎能获得最佳的GPU加速。
- Core ML 或 TFLite：针对苹果设备或安卓设备，分别考虑转换为Core ML或TensorFlow Lite格式。
使用优化过的推理库：考虑使用专门为Whisper优化的推理库，例如faster-whisper（基于CTranslate2），它本身就更高效，再结合我们剪枝后的模型，效果会更好。
内存与速度的权衡：在边缘设备上，内存往往比算力更稀缺。如果你的设备内存非常紧张，可能需要在剪枝时更激进一些（牺牲更多精度来换取更小的模型），或者采用更极端的量化方法。