4. AI模型打包

在Edge AI（边缘人工智能）的工程实践中，训练完成的深度学习模型仅仅是一个“半成品”。边缘设备（如Jetson、RK3588、树莓派、手机NPU等）的算力、内存和功耗都远逊于云端GPU服务器，直接部署原始权重文件往往会导致推理延迟过高、内存溢出甚至无法运行。本章聚焦于模型打包这一承上启下的关键环节，系统讲解如何将PyTorch/TensorFlow模型导出为ONNX或TorchScript中间格式，如何通过量化与剪枝实现模型瘦身，以及如何组织模型文件以适配不同的推理引擎和部署场景。下图展示了从训练框架到边缘推理引擎的完整模型打包链路：

上图清晰地展示了模型打包的三大层次：训练框架层负责产出原始权重并执行优化；中间表示层通过ONNX或TorchScript实现框架解耦；边缘推理层则根据硬件特性选择对应的运行时引擎。整个链路的核心目标是：在保持模型精度的前提下，最大化推理速度并最小化资源占用。

4.1 模型导出：从训练框架到中间表示

4.1.1 核心原理：为什么需要中间表示

PyTorch和TensorFlow各自拥有独立的计算图表示和运行时。直接使用原始权重文件部署会带来两个问题：一是目标设备可能不支持该框架的运行时；二是框架自身的推理优化（如算子融合、内存复用）远不如专用推理引擎。ONNX（Open Neural Network Exchange）作为一种开放格式，将模型转换为与框架无关的计算图，使得同一份模型文件可以被ONNX Runtime、TensorRT、OpenVINO等引擎加载。TorchScript则是PyTorch原生的序列化格式，它通过追踪（Tracing）或脚本（Scripting）将Python动态图固化为静态图，适合在C++环境中无Python依赖地执行。

4.1.2 操作流程：PyTorch导出ONNX与TorchScript

PyTorch导出ONNX的核心API是torch.onnx.export。以下是一个典型的导出流程：

import torch
import torchvision.models as models

# 1. 加载预训练模型并设为评估模式
model = models.resnet18(pretrained=True)
model.eval()

# 2. 构造示例输入（batch=1, 3通道, 224x224）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,                      # 模型
    dummy_input,                # 示例输入
    "resnet18.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"}}
)
print("ONNX导出完成")

导出TorchScript有两种方式：追踪（Tracing）和脚本（Scripting）。追踪适用于不含数据依赖控制流的模型，脚本则能处理if/for等动态逻辑。

# 方式一：Tracing（追踪）
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save("resnet18_traced.pt")

# 方式二：Scripting（脚本）
scripted_model = torch.jit.script(model)
scripted_model.save("resnet18_scripted.pt")

4.1.3 案例拆解：TensorFlow模型导出ONNX

TensorFlow模型导出ONNX通常借助tf2onnx工具。假设我们有一个训练好的Keras模型：

import tensorflow as tf
import tf2onnx

# 加载Keras模型
model = tf.keras.models.load_model("mobilenet_v2.h5")

# 转换为ONNX
onnx_model, _ = tf2onnx.convert.from_keras(model, opset=13, output_path="mobilenet_v2.onnx")
print("转换完成")

转换后务必使用onnxruntime进行验证，确保输出与原始模型一致：

import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("mobilenet_v2.onnx")
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
dummy_input = np.random.rand(1, 224, 224, 3).astype(np.float32)
result = sess.run([output_name], {input_name: dummy_input})
print("ONNX推理输出shape:", result[0].shape)

4.1.4 常见误区与总结

此外，动态轴（dynamic_axes）虽然灵活，但会降低推理引擎的优化空间。如果边缘设备的输入尺寸固定，应尽量使用静态shape以获得最佳性能。

4.2 模型优化：量化与剪枝

4.2.1 核心原理：精度与效率的权衡

边缘设备的算力有限，FP32模型往往无法满足实时性要求。量化（Quantization）通过将权重和激活从FP32降低到INT8或FP16，显著减少模型体积和计算量。剪枝（Pruning）则通过移除不重要的权重连接或通道，直接缩小模型结构。两者可以组合使用，但需要仔细评估精度损失。

优化技术	原理	典型收益	精度影响	适用场景
FP16半精度	权重/激活用16位浮点	体积减半，速度提升1.5-2x	几乎无损	GPU/NPU支持FP16
INT8动态量化	仅权重量化，激活运行时计算	体积减至1/4	1-3%下降	CPU推理
INT8全量化	权重+激活均量化，需校准数据集	体积减至1/4，速度提升3-5x	2-5%下降	TensorRT/OpenVINO
结构化剪枝	移除整个通道/滤波器	体积减少30-50%	需微调恢复	卷积网络
非结构化剪枝	移除单个权重	稀疏度高但硬件加速难	需微调	研究实验

4.2.2 操作流程：PyTorch量化实践

PyTorch提供了三种量化方案：动态量化、静态量化和量化感知训练（QAT）。对于边缘部署，静态量化通常效果最佳。以下是一个静态量化流程：

import torch
from torch.quantization import quantize_fx

# 准备模型
model = models.resnet18(pretrained=True).eval()

# 准备校准数据集（示例：随机数据）
calibration_data = [torch.randn(1, 3, 224, 224) for _ in range(100)]

# 使用FX量化
qconfig = torch.quantization.get_default_qconfig("fbgemm")
q_model = quantize_fx.prepare_fx(model, {"": qconfig}, calibration_data[0])
# 校准
for data in calibration_data:
    q_model(data)
# 转换
q_model = quantize_fx.convert_fx(q_model)
torch.jit.save(torch.jit.script(q_model), "resnet18_int8.pt")
print("INT8量化完成")

4.2.3 案例拆解：剪枝与量化组合优化

假设我们有一个用于工业缺陷检测的YOLOv5s模型，原始FP32模型大小为14.8MB，在Jetson Nano上推理延迟为85ms。我们采用以下优化策略：

1. 通道剪枝：使用torch.nn.utils.prune对C3模块的卷积通道进行L1范数剪枝，剪枝率30%。
2. 微调：在缺陷数据集上微调10个epoch，mAP从0.892恢复至0.881。
3. INT8量化：使用TensorRT的PTQ（训练后量化）进行INT8量化，校准集为500张验证图片。

优化结果对比如下：

模型版本	体积	推理延迟	mAP
FP32原始	14.8MB	85ms	0.892
剪枝后FP32	10.2MB	62ms	0.881
剪枝+INT8	2.8MB	18ms	0.874

最终模型体积缩小81%，延迟降低79%，精度仅损失2%。这一案例表明，剪枝与量化组合使用能带来显著的工程收益。

4.2.4 常见误区与总结

剪枝并非越多越好。过高的剪枝率会导致模型容量不足，即使微调也难以恢复精度。建议从20-30%的剪枝率开始，逐步增加并观察验证集指标。

4.3 模型文件组织与部署规范

4.3.1 核心原理：文件组织的工程意义

在容器化部署场景中，模型文件不仅仅是权重，还包含配置文件、标签映射、预处理参数等。合理的文件组织能简化容器镜像构建、支持版本回滚、便于A/B测试。一个标准的模型目录应包含：模型文件、配置文件、元数据（版本、精度、输入输出格式）、以及可选的示例代码。

4.3.2 操作流程：标准模型目录结构

model_repository/
├── resnet18/
│   ├── 1/                      # 版本号目录
│   │   ├── model.onnx          # 模型文件
│   │   └── config.json         # 模型配置
│   ├── 2/                      # 第二版本
│   │   ├── model.onnx
│   │   └── config.json
│   ├── labels.txt              # 类别标签
│   ├── preprocess.py           # 预处理脚本（可选）
│   └── README.md               # 模型说明
└── yolov5s/
    ├── 1/
    │   ├── model.pt            # TorchScript模型
    │   └── config.yaml         # 超参数配置
    └── metadata.json           # 模型元数据

对于ONNX模型，建议同时保存输入输出的shape、dtype、归一化参数等信息到JSON文件，方便推理服务动态加载。

4.3.3 案例拆解：Triton Inference Server模型仓库

NVIDIA Triton Inference Server要求模型仓库遵循特定目录结构。以下是一个完整的ONNX模型仓库配置：

model_repository/
└── densenet_onnx/
    ├── 1/
    │   └── model.onnx
    └── config.pbtxt

config.pbtxt内容如下：

name: "densenet_onnx"
platform: "onnxruntime_onnx"
max_batch_size: 8
input [
  {
    name: "data_input"
    data_type: TYPE_FP32
    dims: [3, 224, 224]
  }
]
output [
  {
    name: "prob_output"
    data_type: TYPE_FP32
    dims: [1000]
  }
]
dynamic_batching { preferred_batch_size: [1, 4, 8] }

这种组织方式使得Triton能够自动管理模型版本、动态批处理和并发推理，极大简化了边缘服务器的部署复杂度。

4.3.4 常见误区与总结

模型文件组织看似简单，却是DevOps实践中容易忽视的环节。良好的组织规范能显著提升模型迭代效率，降低运维事故风险。

4.4 模型验证与性能基准测试

4.4.1 核心原理：验证是部署的最后一道防线

模型导出和优化后，必须经过严格的验证才能进入生产环境。验证包括两个层面：功能正确性（输出与原始模型一致）和性能达标性（延迟、吞吐量满足SLA）。

4.4.2 操作流程：基准测试方法

import onnxruntime as ort
import time
import numpy as np

sess = ort.InferenceSession("model_int8.onnx", providers=["TensorrtExecutionProvider", "CUDAExecutionProvider"])
input_name = sess.get_inputs()[0].name
dummy_input = np.random.rand(1, 3, 224, 224).astype(np.float32)

# 预热
for _ in range(10):
    sess.run(None, {input_name: dummy_input})

# 正式测试
latencies = []
for _ in range(100):
    start = time.perf_counter()
    sess.run(None, {input_name: dummy_input})
    latencies.append((time.perf_counter() - start) * 1000)

latencies.sort()
p50 = latencies[50]
p95 = latencies[95]
p99 = latencies[99]
print(f"P50: {p50:.2f}ms, P95: {p95:.2f}ms, P99: {p99:.2f}ms")

4.4.3 案例拆解：Jetson Nano上的性能对比

我们在Jetson Nano（Maxwell架构，128 CUDA核心）上对ResNet18的三种版本进行基准测试：

模型版本	推理引擎	P50延迟	P95延迟	吞吐量
FP32 ONNX	ONNX Runtime CPU	42.3ms	45.1ms	23.6 FPS
FP16 ONNX	TensorRT FP16	18.7ms	20.2ms	53.5 FPS
INT8 ONNX	TensorRT INT8	9.8ms	10.5ms	102.0 FPS

测试结果表明，INT8量化在Jetson Nano上实现了4.3倍的加速，完全满足实时视频流分析（30 FPS）的需求。

4.4.4 常见误区与总结

模型验证不是一次性工作。当推理引擎版本升级、硬件驱动更新或输入数据分布变化时，都应重新执行基准测试，确保性能指标仍然达标。

本章系统讲解了AI模型打包的完整流程：从PyTorch/TensorFlow导出ONNX或TorchScript，到量化与剪枝优化，再到模型文件组织与验证。掌握这些技术，能够帮助开发者将训练好的模型高效、稳定地部署到各类边缘设备上，为后续的容器化推理服务搭建奠定坚实基础。