深度解析：8位量化对GPU上Whisper模型推理速度的影响及应用场景

import torch
from transformers import WhisperFeatureExtractor, WhisperTokenizerFast, AutoModelForSpeechSeq2Seq
from transformers.pipelines.audio_classification import ffmpeg_read
import time # 引入time模块用于计时

# 模型名称
MODEL_NAME = "openai/whisper-large-v3"

# 初始化特征提取器和分词器
tokenizer = WhisperTokenizerFast.from_pretrained(MODEL_NAME)
feature_extractor = WhisperFeatureExtractor.from_pretrained(MODEL_NAME)

print("加载8位量化模型...")
# 加载8位量化模型
# device_map='auto' 会自动将模型分配到可用设备，load_in_8bit=True 启用8位量化
model_8bit = AutoModelForSpeechSeq2Seq.from_pretrained(
    MODEL_NAME,
    device_map='auto',
    load_in_8bit=True
)
print("8位量化模型加载完成。")

# 示例音频文件路径 (请替换为您的实际文件)
sample_file = "sample.mp3" # 假设这是一个27秒长的音频文件

print(f"开始对文件 '{sample_file}' 进行推理...")
with torch.inference_mode(): # 禁用梯度计算，节省内存并加速推理
    with open(sample_file, "rb") as f:
        inputs_bytes = f.read()
        # 使用ffmpeg_read读取音频数据并重采样
        audio_data = ffmpeg_read(inputs_bytes, feature_extractor.sampling_rate)

        # 提取音频特征
        input_features = feature_extractor(
            audio_data,
            sampling_rate=feature_extractor.sampling_rate,
            return_tensors='pt'
        )['input_features']

        # 将输入特征移动到GPU并转换为float16精度
        # 注意：即使模型是8位，输入特征通常仍以更高精度处理
        input_features = input_features.to(dtype=torch.float16, device='cuda')

        start_time = time.time() # 记录推理开始时间

        # 进行模型推理
        forced_decoder_ids_output = model_8bit.generate(
            input_features=input_features,
            return_timestamps=False # 根据需求设置是否返回时间戳
        )

        end_time = time.time() # 记录推理结束时间
        inference_time = end_time - start_time
        print(f"推理完成，耗时: {inference_time:.2f} 秒")

        # 解码输出
        out_text = tokenizer.decode(forced_decoder_ids_output.squeeze())
        print("识别结果:", out_text)