浙江市建设网站,共享设计平台,西安网页制作与设计,佛山建设Whisper-large-v3模型剪枝教程#xff1a;减小模型大小保持精度 你是不是也遇到过这种情况#xff1a;想把一个强大的语音识别模型#xff0c;比如Whisper-large-v3#xff0c;塞进你的边缘设备里#xff0c;结果发现它太大了#xff0c;根本装不下#xff1f;或者就算…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 PromptWindows执行以下命令# 创建一个新的Python环境命名为whisper_prune conda create -n whisper_prune python3.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, modelmodel_id, devicedevice, chunk_length_s30 # 处理长音频时切片 ) # 2. 加载一个小型测试数据集这里用Distil-Whisper提供的一个干净样本 print(加载测试数据...) dataset load_dataset(distil-whisper/librispeech_long, clean, splitvalidation) # 我们只取前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样本 {i1}:) 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_inputexample_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(dim1) # 对行求和 # 假设我们想剪掉该层50%的通道 prune_ratio 0.5 num_prune int(len(importance) * prune_ratio) # 找到最不重要的通道索引 _, idxs torch.topk(importance, knum_prune, largestFalse) 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, splittrain.100) # 只用100个样本演示 # 预处理函数将音频转换为模型输入格式 def prepare_dataset(batch): audio batch[audio] # 重采样到16kHzWhisper的标准输入 if audio[sampling_rate] ! 16000: # 这里需要音频处理库如torchaudio或librosa # 为简化假设数据已是16kHz pass batch[input_features] processor(audio[array], sampling_rateaudio[sampling_rate], return_tensorspt).input_features[0] # 处理文本标签 batch[labels] processor.tokenizer(batch[text]).input_ids return batch # 应用预处理 dataset dataset.map(prepare_dataset, remove_columnsdataset.column_names) # 3. 分割训练集和验证集 split_dataset dataset.train_test_split(test_size0.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_tensorspt) # 填充标签 labels_batch self.processor.tokenizer.pad(label_features, return_tensorspt) # 将填充的标签token替换为-100以便在损失计算时忽略 labels labels_batch[input_ids].masked_fill(labels_batch.attention_mask.ne(1), -100) batch[labels] labels return batch data_collator DataCollatorSpeechSeq2SeqWithPadding(processorprocessor) # 5. 配置训练参数关键步骤 # 对于剪枝后的模型学习率可以设小一点训练轮数epoch也不用太多 training_args Seq2SeqTrainingArguments( output_dir./whisper_pruned_finetuned, per_device_train_batch_size2, # 根据你的GPU内存调整 per_device_eval_batch_size2, gradient_accumulation_steps4, # 模拟更大的批次 learning_rate1e-5, # 较小的学习率 warmup_steps50, num_train_epochs3, # 微调3轮通常足够 logging_dir./logs, logging_steps10, evaluation_strategyepoch, save_strategyepoch, load_best_model_at_endTrue, metric_for_best_modelwer, # 根据WER选择最佳模型 greater_is_betterFalse, push_to_hubFalse, # 设为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_tokensTrue) label_ids[label_ids -100] processor.tokenizer.pad_token_id label_str processor.batch_decode(label_ids, skip_special_tokensTrue) # 计算WER wer 100 * metric.compute(predictionspred_str, referenceslabel_str) return {wer: wer} # 7. 创建Trainer并开始微调 trainer Seq2SeqTrainer( modelmodel, argstraining_args, train_datasettrain_dataset, eval_dataseteval_dataset, data_collatordata_collator, compute_metricscompute_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, devicecuda): 评估指定模型的WER、推理速度和大小 # 加载模型和管道 print(f\n评估模型: {model_path}) pipe pipeline( automatic-speech-recognition, modelmodel_path, devicedevice, torch_dtypetorch.float16 if devicecuda 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 样本{i1}识别: {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, splitvalidation) 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}, dtypetorch.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它本身就更高效再结合我们剪枝后的模型效果会更好。内存与速度的权衡在边缘设备上内存往往比算力更稀缺。如果你的设备内存非常紧张可能需要在剪枝时更激进一些牺牲更多精度来换取更小的模型或者采用更极端的量化方法。走完这一整套流程你应该已经得到了一个比原版Whisper-large-v3小得多、也快得多的模型并且识别精度仍然保持在可用的水平。这个过程可能需要一些迭代和调优比如调整剪枝比例、尝试不同的微调策略等但基本框架就是如此。最重要的是你现在有了一个可以在资源受限环境中实际运行的语音识别能力这为开发各种离线语音应用打开了大门。无论是做智能录音笔、会议转录盒子还是嵌入到机器人或IoT设备中这个“瘦身成功”的Whisper都能派上用场。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。