加载 GPT-2 量化模型并推理输出：INT8 推理流程与输出解析

本文基于 quant_infer.py 脚本，讲解如何加载动态量化后的 GPT-2 小模型并完成 prompt 推理，展示 from_pretrained + load_state_dict 的加载技巧，封装输出预测函数 infer(prompt)，实现完整的 INT8 推理流程。

---

📦 模型准备说明

在第一篇文章中，我们通过动态量化保存了如下模型：

gpt2_student_v2_quantized/
├── pytorch_model.bin       # INT8 权重
├── config.json             # 模型结构配置

模型结构保持不变，只是 Linear 层的参数被压缩为 INT8 格式。

---

1️⃣ 加载量化模型（INT8 权重）

import torch
from transformers import GPT2LMHeadModel

model_path = "./gpt2_student_v2_quantized"

# 1. 初始化空模型结构
model = GPT2LMHeadModel.from_pretrained(model_path)

# 2. 加载量化权重
state_dict = torch.load(f"{model_path}/pytorch_model.bin")
model.load_state_dict(state_dict)

# 3. 设置 eval 模式
model.eval()

⚠️ 说明：

• from_pretrained() 会自动读取 config.json
• load_state_dict() 加载的是量化参数
• 必须 model.eval() 否则推理结果不稳定

---

2️⃣ 加载 tokenizer（原始路径）

from transformers import GPT2Tokenizer

tokenizer = GPT2Tokenizer.from_pretrained("../python3_distillation/gpt2_student_v2")

量化模型与 tokenizer 没有冲突，词表结构不变。

---

3️⃣ 封装推理函数 infer(prompt)

def infer(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    logits = outputs.logits
    top1_token_id = int(logits[0, -1].argmax())
    token = tokenizer.decode([top1_token_id], clean_up_tokenization_spaces=True).strip()

    return {
        "prompt": prompt,
        "token_id": top1_token_id,
        "token": token
    }

---

4️⃣ 多 prompt 测试示例

prompts = [
    "你好，今天是",
    "I love",
    "Artificial intelligence is",
    "GPT-2 is a",
]

for prompt in prompts:
    result = infer(prompt)
    print(f"📝 Prompt: {result['prompt']}")
    print(f"🔹 Token ID: {result['token_id']}")
    print(f"🔹 Token: {result['token']}")

输出示例：

📝 Prompt: I love
🔹 Token: you

📝 Prompt: GPT-2 is a
🔹 Token: powerful

---

5️⃣ 输出结构说明

• logits shape 为 [B, L, V]
• logits[:, -1, :] 取最后一个 token 的预测结果
• argmax(dim=-1) 得到最可能的下一个 token id
• tokenizer.decode() 转为字符串输出

---

✅ 总结

• 量化模型加载需配合 from_pretrained + load_state_dict
• infer(prompt) 函数支持批量封装与 API 接入
• 推理输出与原始模型基本一致，速度明显提升
• 推荐用于 CPU 本地服务、嵌入式部署、轻量级任务调用

---

🧭 本系列 GPT-2 模型量化部署系列四部曲

• 🧩 第一篇：GPT-2 模型动态量化实战：一键压缩为 INT8 权重，提升部署效率
• 🚀 第二篇：加载 GPT-2 量化模型并推理输出：INT8 推理流程与输出解析战
• 🌐 第三篇：GPT-2 模型量化前后推理性能对比：速度提升与输出一致性实测
• 🧠 第四篇：用 Flask 封装 GPT-2 INT8 推理服务：轻量部署与接口调用实战

---

📌 YoanAILab 技术导航页

💡 项目源码 × 实战部署 × 转型经验，一页总览
👉 点击查看完整导航页

📚 包含内容：

• 🧠 GPT-2 项目源码（GitHub）
• ✍️ CSDN 技术专栏合集
• 💼 知乎转型日志
• 📖 公众号 YoanAILab 全文合集