边缘 AI 量化部署：DMS 模型 INT8 优化实战

量化背景

边缘部署挑战：

指标	云端 GPU	车载 NPU
内存	16-32 GB	2-4 GB
算力	100+ TOPS	10-30 TOPS
功耗	300W	< 10W
精度	FP32/FP16	INT8/INT4

量化目标：

• 模型大小减小 4x
• 推理速度提升 2-3x
• 精度损失 < 1%

---

量化方法

PTQ vs QAT

方法	描述	优势	劣势
PTQ	训练后量化	快速、简单	精度损失大
QAT	量化感知训练	精度高	训练复杂

INT8 量化实现

"""
DMS 模型 INT8 量化

支持框架：
1. PyTorch -> ONNX -> TensorRT
2. PyTorch -> TFLite
3. ONNX Runtime INT8
"""

import torch
import torch.nn as nn
import onnx
import onnxruntime as ort
import numpy as np

class DMSModel(nn.Module):
    """简化 DMS 模型"""
    
    def __init__(self):
        super().__init__()
        
        # 人脸特征提取
        self.face_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            
            nn.Conv2d(64, 128, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1)
        )
        
        # 关键点回归
        self.keypoint_head = nn.Linear(128, 68 * 2)  # 68 个关键点
        
        # 疲劳分类
        self.fatigue_head = nn.Linear(128, 2)
        
        # 视线估计
        self.gaze_head = nn.Linear(128, 2)
    
    def forward(self, x):
        # 特征提取
        feat = self.face_encoder(x)
        feat = feat.view(feat.size(0), -1)
        
        # 多任务输出
        keypoints = self.keypoint_head(feat)
        fatigue = self.fatigue_head(feat)
        gaze = self.gaze_head(feat)
        
        return {
            "keypoints": keypoints,
            "fatigue": fatigue,
            "gaze": gaze
        }


def export_to_onnx(model, output_path, input_shape=(1, 3, 224, 224)):
    """导出 ONNX"""
    model.eval()
    
    dummy_input = torch.randn(*input_shape)
    
    torch.onnx.export(
        model,
        dummy_input,
        output_path,
        input_names=["input"],
        output_names=["keypoints", "fatigue", "gaze"],
        dynamic_axes={
            "input": {0: "batch"},
            "keypoints": {0: "batch"},
            "fatigue": {0: "batch"},
            "gaze": {0: "batch"}
        },
        opset_version=11
    )
    
    print(f"ONNX 模型已导出: {output_path}")
    return output_path


def quantize_onnx_ptq(onnx_path, calibration_data, output_path):
    """
    ONNX PTQ 量化
    
    Args:
        onnx_path: FP32 ONNX 模型路径
        calibration_data: 校准数据
        output_path: 输出路径
    """
    from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantFormat, QuantType
    from onnxruntime.quantization.shape_inference import quant_pre_process
    
    # 预处理
    preprocessed_path = onnx_path.replace(".onnx", "_preprocessed.onnx")
    quant_pre_process(onnx_path, preprocessed_path)
    
    # 静态量化（需要校准数据）
    quantize_static(
        preprocessed_path,
        output_path,
        calibration_data_reader=calibration_data,
        quant_format=QuantFormat.QDQ,
        per_channel=False,
        weight_type=QuantType.QInt8
    )
    
    print(f"INT8 模型已保存: {output_path}")
    
    return output_path


def benchmark_model(model_path, input_shape=(1, 3, 224, 224), num_runs=100):
    """性能基准测试"""
    import time
    
    session = ort.InferenceSession(model_path)
    input_name = session.get_inputs()[0].name
    
    # 准备数据
    input_data = np.random.randn(*input_shape).astype(np.float32)
    
    # 预热
    for _ in range(10):
        session.run(None, {input_name: input_data})
    
    # 计时
    start = time.time()
    for _ in range(num_runs):
        session.run(None, {input_name: input_data})
    elapsed = time.time() - start
    
    avg_time = elapsed / num_runs * 1000
    fps = 1000 / avg_time
    
    # 内存占用
    model_size = onnx.load(model_path).ByteSize() / 1024 / 1024
    
    return {
        "avg_inference_time_ms": avg_time,
        "fps": fps,
        "model_size_mb": model_size
    }


# 测试流程
if __name__ == "__main__":
    # 1. 创建模型
    model = DMSModel()
    
    # 2. 导出 ONNX
    onnx_path = "dms_model.onnx"
    export_to_onnx(model, onnx_path)
    
    # 3. 量化（简化）
    int8_path = "dms_model_int8.onnx"
    # quantize_onnx_ptq(onnx_path, calibration_data, int8_path)
    
    # 4. 性能对比
    fp32_perf = benchmark_model(onnx_path)
    print(f"FP32 性能: {fp32_perf}")
    
    # int8_perf = benchmark_model(int8_path)
    # print(f"INT8 性能: {int8_perf}")

---

TensorRT 部署

"""
TensorRT INT8 部署

NVIDIA GPU 最优方案
"""

class TensorRTDeployer:
    """TensorRT 部署器"""
    
    def __init__(self):
        self.engine = None
        
    def build_engine(self, onnx_path, batch_size=1):
        """构建 TensorRT 引擎"""
        import tensorrt as trt
        import pycuda.driver as cuda
        
        TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
        
        with trt.Builder(TRT_LOGGER) as builder:
            with builder.create_network(1) as network:
                with trt.OnnxParser(network, TRT_LOGGER) as parser:
                    # 解析 ONNX
                    with open(onnx_path, 'rb') as f:
                        parser.parse(f.read())
                    
                    # 配置
                    config = builder.create_builder_config()
                    config.max_workspace_size = 1 << 30  # 1GB
                    
                    # INT8 模式
                    config.set_flag(trt.BuilderFlag.INT8)
                    
                    # 校准器（INT8 必须）
                    # calibrator = ...
                    # config.int8_calibrator = calibrator
                    
                    # 构建引擎
                    self.engine = builder.build_engine(network, config)
        
        return self.engine
    
    def infer(self, input_data):
        """推理"""
        import pycuda.driver as cuda
        import pycuda.autoinit
        
        context = self.engine.create_execution_context()
        
        # 分配内存
        # ... 简化实现
        
        return {"output": None}

---

精度评估

"""
量化精度评估

对比 FP32 和 INT8 输出差异
"""

def evaluate_quantization_accuracy(fp32_model, int8_model, test_data):
    """
    评估量化精度
    
    Args:
        fp32_model: FP32 模型
        int8_model: INT8 模型
        test_data: 测试数据
        
    Returns:
        精度指标
    """
    import numpy as np
    
    fp32_session = ort.InferenceSession(fp32_model)
    int8_session = ort.InferenceSession(int8_model)
    
    results = []
    
    for batch in test_data:
        # FP32 推理
        fp32_out = fp32_session.run(None, {"input": batch})[0]
        
        # INT8 推理
        int8_out = int8_session.run(None, {"input": batch})[0]
        
        # 计算误差
        mse = np.mean((fp32_out - int8_out) ** 2)
        mae = np.mean(np.abs(fp32_out - int8_out))
        
        results.append({"mse": mse, "mae": mae})
    
    # 统计
    avg_mse = np.mean([r["mse"] for r in results])
    avg_mae = np.mean([r["mae"] for r in results])
    
    return {
        "avg_mse": avg_mse,
        "avg_mae": avg_mae,
        "relative_error_pct": avg_mae * 100,
        "acceptable": avg_mae < 0.01  # 1% 相对误差
    }

---

总结

边缘 AI 量化关键要点：

1. PTQ 快速部署：适合大多数场景
2. QAT 高精度：对精度要求高时使用
3. INT8 平衡点：精度损失小，速度提升明显
4. TensorRT 最优：NVIDIA GPU 最佳性能
5. 验证必须：确保精度在可接受范围

---

参考资料：

1. ONNX Runtime Quantization
2. TensorRT Developer Guide
3. PyTorch Quantization Tutorial