如何将Hugging Face预训练模型转为.pt并本地运行？解决加载报错

错误原因 #

你遇到的报错是因为AutoModelForCTC.from_pretrained()并不支持直接加载单独的.pt权重文件，该方法的参数需要是Hugging Face模型仓库ID，或者包含模型配置文件、权重文件等完整文件的本地文件夹路径。之前只保存了模型的权重字典（state_dict），没有保存模型的结构配置信息，所以加载时无法识别。

正确的模型导出代码 #

直接使用save_pretrained()方法保存完整模型和处理器，会自动生成所有必要的配置文件：

from transformers import AutoModelForCTC, Wav2Vec2Processor

# 定义模型名称
model_name = "UrukHan/wav2vec2-russian"

# 加载模型和处理器
model = AutoModelForCTC.from_pretrained(model_name)
processor = Wav2Vec2Processor.from_pretrained(model_name)

# 将完整模型和处理器保存到本地文件夹
model.save_pretrained("./russian-wav2vec2-local")
processor.save_pretrained("./russian-wav2vec2-local")

正确的本地推理代码 #

从本地文件夹加载模型和处理器，规范完成推理流程：

from transformers import AutoModelForCTC, Wav2Vec2Processor
import librosa
import torch

# 从本地文件夹加载模型和处理器
model = AutoModelForCTC.from_pretrained("./russian-wav2vec2-local")
processor = Wav2Vec2Processor.from_pretrained("./russian-wav2vec2-local")

# 设置模型为评估模式
model.eval()

# 加载音频文件
FILE = "wav.wav"
audio, sr = librosa.load(FILE, sr=16000)

# 用处理器规范预处理音频，生成模型输入张量
input_values = processor(audio, sampling_rate=sr, return_tensors="pt").input_values

# 执行推理
with torch.no_grad():
    logits = model(input_values).logits

# 解码得到转录文本
pred_ids = torch.argmax(logits, dim=-1)
transcription = processor.batch_decode(pred_ids)[0]

print(transcription)

内容的提问来源于stack exchange，提问作者vonexel