3. 模型格式与转换

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

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

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()  # 别忘了切到推理模式

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")

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"}
    }
)

导出后一定要验证一下。我习惯用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()}")

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)

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版本
# 很多模型已经有人转换好了

格式选型：一张表说清楚

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

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