模型格式与转换:从PyTorch到GGUF的实战之路

模型格式这个话题,说实话刚入行时我也觉得挺枯燥的。直到有一次,我在一个边缘设备项目里折腾了整整三天,就为了把一个PyTorch模型塞进手机端跑起来。嗯,从那以后我再也不敢小看格式转换了。

今天咱们就把这几种主流格式捋一遍。你想想看,一个模型从训练到部署,中间要经过多少道工序?格式转换就是那个承上启下的关键环节。

核心认知:模型格式不是简单的文件后缀名,它决定了推理引擎能做什么优化、能跑多快、能部署到哪。选错格式,后面全白干。

模型格式转换全景图 PyTorch (.pt/.pth) 训练态 / 动态图 ONNX (.onnx) 中间表示 / 静态图 TensorRT (.engine) NVIDIA GPU 专属 GGUF (.gguf) CPU / 边缘设备 HuggingFace Hub 模型仓库 / 分发 转换路径 PyTorch → ONNX → 目标格式

PyTorch模型格式:训练态的起点

PyTorch的模型保存,说白了就两种方式。我个人习惯用state_dict,因为它最灵活。

# 方式一:保存完整模型(不推荐用于部署)
torch.save(model, "model_complete.pt")

# 方式二:保存 state_dict(推荐)
torch.save(model.state_dict(), "model_weights.pth")

# 加载时重建
model = MyModel()
model.load_state_dict(torch.load("model_weights.pth"))
model.eval()  # 别忘了切到推理模式

我曾经踩过的坑:直接保存完整模型,换了个Python版本就加载不上了。state_dict只存参数,跨版本兼容性好得多。另外,model.eval()这行我忘过不止一次,dropout和batch norm在推理时乱跳,结果对不上。

HuggingFace模型仓库:生态的枢纽

HuggingFace Hub现在几乎是模型分发的标准了。你想想看,从Llama到Qwen,从BERT到Whisper,全都在上面。我个人建议,不管模型多大,都走from_pretrained接口。

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")

# 保存到本地
model.save_pretrained("./local_qwen")
tokenizer.save_pretrained("./local_qwen")

小技巧:HuggingFace的模型默认包含config.json和tokenizer文件。我建议你下载后先检查model.safetensors.index.json,确认分片信息完整。有一次我漏了某个分片,加载时卡了半小时才报错。

ONNX导出:打通部署的桥梁

ONNX是个好东西。它把PyTorch的动态图转成静态图,说白了就是给模型画了一张固定的计算流程图。这样下游的TensorRT、ONNX Runtime才能做各种优化。

import torch
import torch.onnx

# 准备一个示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,                     # PyTorch模型
    dummy_input,               # 示例输入
    "model.onnx",              # 输出路径
    export_params=True,        # 导出参数
    opset_version=17,          # ONNX算子集版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=["input"],     # 输入名称
    output_names=["output"],   # 输出名称
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)

关键参数说明:

  • opset_version:越高支持的算子越多,但兼容性可能下降。我一般用17或18。
  • dynamic_axes:允许动态batch size。如果你固定batch size推理,可以不加。
  • do_constant_folding:把能算的常量提前算好,减少推理时的计算量。

导出后一定要验证一下。我习惯用ONNX Runtime跑一遍,对比输出差异。

import onnx
import onnxruntime as ort

# 验证ONNX模型结构
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)

# 用ONNX Runtime推理
session = ort.InferenceSession("model.onnx")
inputs = {session.get_inputs()[0].name: dummy_input.numpy()}
outputs = session.run(None, inputs)

# 对比PyTorch输出
with torch.no_grad():
    torch_output = model(dummy_input)
print(f"最大差异: {torch.abs(torch_output - torch.tensor(outputs[0])).max()}")

我曾经遇到的坑:ONNX导出时,如果模型里有自定义算子,会直接报错。解决办法要么是注册自定义算子,要么是改模型结构避开。还有一次,我忘了设dynamic_axes,结果生产环境来了个不同batch size的请求,直接崩了。

TensorRT模型转换:GPU上的极致优化

TensorRT是NVIDIA的推理优化神器。它能把ONNX模型编译成GPU上跑得飞快的engine文件。说白了,就是给模型做了一次深度定制优化。

import tensorrt as trt

# 创建builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

# 解析ONNX模型
with open("model.onnx", "rb") as f:
    parser.parse(f.read())

# 配置优化
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
config.set_flag(trt.BuilderFlag.FP16)  # 启用FP16

# 构建engine
serialized_engine = builder.build_serialized_network(network, config)
with open("model.engine", "wb") as f:
    f.write(serialized_engine)

实战建议:TensorRT的编译时间可能很长,大模型动辄半小时。我建议你在CI/CD里提前编译好,把engine文件作为制品保存。另外,FP16精度损失通常可以忽略,但显存占用直接减半,性价比极高。

GGUF格式:CPU推理的救星

GGUF是llama.cpp带起来的格式,专门为CPU和边缘设备设计。它把模型量化、内存映射、KV缓存全打包在一起。说实话,我第一次在MacBook上跑7B模型时,真被GGUF的流畅度惊到了。

GGUF的核心优势在于:

  • 内存映射加载:不需要把整个模型读进内存,按需加载
  • 多种量化方式:从Q2到Q8,精度和速度自由权衡
  • 单文件分发:一个.gguf文件包含所有东西,部署超简单
# 使用llama.cpp的convert脚本转换
python convert.py ./local_qwen \
    --outfile qwen2.5-7b-q4_k_m.gguf \
    --outtype q4_k_m

# 或者用HuggingFace直接下载GGUF版本
# 很多模型已经有人转换好了

量化级别选择指南:

量化类型 每参数位数 推荐场景
Q2_K 2.56 极致压缩,质量损失明显
Q4_K_M 4.50 最佳性价比,我常用这个
Q5_K_M 5.06 质量接近FP16,体积适中
Q8_0 8.00 几乎无损,体积较大

我曾经犯过的错:第一次用GGUF时,我直接选了Q2_K想压到最小,结果模型输出全是乱码。后来才明白,量化不是越低越好,得根据任务来。对话任务至少Q4,代码生成建议Q5以上。

格式选型:一张表说清楚

格式 适用场景 优化程度 部署难度
PyTorch (.pt/.pth) 训练、调试、实验 无 低
ONNX (.onnx) 跨平台中间表示 中等 中
TensorRT (.engine) NVIDIA GPU 生产部署 极高 高
GGUF (.gguf) CPU / 边缘 / 移动设备 高(量化) 低
HuggingFace Hub 模型分发、版本管理 取决于源格式 低

选型其实没那么复杂。我个人习惯是:训练用PyTorch,跨平台用ONNX,NVIDIA GPU用TensorRT,CPU或边缘设备用GGUF。HuggingFace Hub作为分发中心,贯穿始终。

最后一个小建议:不管选什么格式,一定要在目标硬件上做端到端测试。我见过太多人格式转换成功了,但推理结果对不上。写个自动化测试脚本,每次转换完自动对比输出,能省很多排查时间。