DMS 模型边缘优化实战：TensorRT + ONNX Runtime 量化部署

发表于 更新于 分类于

发布时间： 2026-04-14
关键词： Model Optimization、TensorRT、ONNX Runtime、Quantization、Edge AI

模型优化流程

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
┌─────────────────────────────────────────────────────┐
│           DMS 模型优化流程                          │
├─────────────────────────────────────────────────────┤
│                                                     │
│   训练模型 (PyTorch/TensorFlow)                    │
│   ─────────────────────────                         │
│       │                                             │
│       ▼                                             │
│   ┌─────────────────┐                              │
│   │ 导出 ONNX       │  标准化中间格式              │
│   └────────┬────────┘                              │
│            │                                        │
│            ▼                                        │
│   ┌─────────────────────────────────────────┐      │
│   │           模型优化                       │      │
│   │  • 剪枝 (Pruning)                       │      │
│   │  • 量化 (Quantization)                  │      │
│   │  • 知识蒸馏 (Distillation)              │      │
│   └────────┬────────────────────────────────┘      │
│            │                                        │
│            ▼                                        │
│   ┌─────────────────┐                              │
│   │ 部署引擎        │  TensorRT / ONNX Runtime    │
│   └────────┬────────┘                              │
│            │                                        │
│            ▼                                        │
│   ┌─────────────────┐                              │
│   │ NPU/GPU 推理    │  车载 SoC                   │
│   └─────────────────┘                              │
│                                                     │
└─────────────────────────────────────────────────────┘

量化技术详解

INT8 量化原理

精度 存储空间 计算速度 准确率损失
FP32 4 字节 基准 0%
FP16 2 字节 ~2x < 0.5%
INT8 1 字节 ~4x 1-3%

量化代码实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
import torch
import torch.nn as nn
import onnx
from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantType
from onnxruntime.quantization.shape_inference import quant_pre_process

class DMSModelQuantizer:
    """DMS 模型量化器"""
    
    def __init__(self, pytorch_model_path: str, output_dir: str):
        self.pytorch_model_path = pytorch_model_path
        self.output_dir = output_dir
        
    def export_to_onnx(self, 
                       model: nn.Module,
                       input_shape: tuple = (1, 3, 224, 224),
                       opset_version: int = 14):
        """导出 ONNX 模型
        
        Args:
            model: PyTorch 模型
            input_shape: 输入形状
            opset_version: ONNX opset 版本
        """
        model.eval()
        
        # 创建示例输入
        dummy_input = torch.randn(*input_shape)
        
        # 导出
        onnx_path = f"{self.output_dir}/dms_model.onnx"
        torch.onnx.export(
            model,
            dummy_input,
            onnx_path,
            opset_version=opset_version,
            input_names=['input'],
            output_names=['output'],
            dynamic_axes={
                'input': {0: 'batch_size'},
                'output': {0: 'batch_size'}
            }
        )
        
        print(f"ONNX model exported to: {onnx_path}")
        return onnx_path
    
    def dynamic_quantization(self, onnx_path: str):
        """动态量化
        
        优点：无需校准数据
        缺点：仅支持部分算子
        """
        quantized_path = onnx_path.replace('.onnx', '_dynamic_int8.onnx')
        
        quantize_dynamic(
            model_input=onnx_path,
            model_output=quantized_path,
            weight_type=QuantType.QInt8,
            per_channel=True,  # 每通道量化
            reduce_range=True
        )
        
        print(f"Dynamic quantized model: {quantized_path}")
        return quantized_path
    
    def static_quantization(self, 
                            onnx_path: str,
                            calibration_data: list):
        """静态量化
        
        优点：精度更高，推理更快
        缺点：需要校准数据
        """
        # 预处理
        preprocessed_path = onnx_path.replace('.onnx', '_preprocessed.onnx')
        quant_pre_process(onnx_path, preprocessed_path, skip_symbolic_shape=True)
        
        # 创建校准数据读取器
        from onnxruntime.quantization import CalibrationDataReader
        
        class DMSDataReader(CalibrationDataReader):
            def __init__(self, calibration_data):
                self.data = calibration_data
                self.index = 0
                
            def get_next(self):
                if self.index >= len(self.data):
                    return None
                data = self.data[self.index]
                self.index += 1
                return {'input': data}
            
            def rewind(self):
                self.index = 0
        
        dr = DMSDataReader(calibration_data)
        
        # 静态量化
        quantized_path = onnx_path.replace('.onnx', '_static_int8.onnx')
        
        quantize_static(
            model_input=preprocessed_path,
            model_output=quantized_path,
            calibration_data_reader=dr,
            quant_format=QuantFormat.QDQ,  # Quantize-Dequantize 格式
            per_channel=True,
            weight_type=QuantType.QInt8,
            activation_type=QuantType.QUInt8
        )
        
        print(f"Static quantized model: {quantized_path}")
        return quantized_path


class TensorRTOptimizer:
    """TensorRT 优化器"""
    
    def __init__(self, onnx_path: str, engine_path: str):
        self.onnx_path = onnx_path
        self.engine_path = engine_path
        
    def build_engine(self, 
                     precision: str = 'fp16',
                     max_batch_size: int = 1,
                     max_workspace_size: int = 1 << 30):  # 1GB
        """构建 TensorRT 引擎
        
        Args:
            precision: 'fp32' | 'fp16' | 'int8'
            max_batch_size: 最大批次大小
            max_workspace_size: 最大工作空间大小
        """
        import tensorrt as trt
        
        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(self.onnx_path, 'rb') as f:
            if not parser.parse(f.read()):
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                return None
        
        # 配置构建器
        config = builder.create_builder_config()
        config.max_workspace_size = max_workspace_size
        
        # 设置精度
        if precision == 'fp16':
            config.set_flag(trt.BuilderFlag.FP16)
        elif precision == 'int8':
            config.set_flag(trt.BuilderFlag.INT8)
            # 需要设置 INT8 校准器
            # config.int8_calibrator = ...
        
        # 构建引擎
        engine = builder.build_engine(network, config)
        
        # 保存引擎
        with open(self.engine_path, 'wb') as f:
            f.write(engine.serialize())
        
        print(f"TensorRT engine saved to: {self.engine_path}")
        return engine
    
    def inference(self, input_data: np.ndarray) -> np.ndarray:
        """TensorRT 推理"""
        import tensorrt as trt
        import pycuda.driver as cuda
        import pycuda.autoinit
        
        # 加载引擎
        with open(self.engine_path, 'rb') as f:
            engine = trt.Runtime(trt.Logger(trt.Logger.WARNING)).deserialize_cuda_engine(f.read())
        
        context = engine.create_execution_context()
        
        # 分配内存
        input_shape = engine.get_binding_shape(0)
        output_shape = engine.get_binding_shape(1)
        
        input_size = trt.volume(input_shape) * np.dtype(np.float32).itemsize
        output_size = trt.volume(output_shape) * np.dtype(np.float32).itemsize
        
        # 分配 GPU 内存
        d_input = cuda.mem_alloc(input_size)
        d_output = cuda.mem_alloc(output_size)
        
        # 分配 CPU 内存
        h_input = input_data.astype(np.float32)
        h_output = np.empty(output_shape, dtype=np.float32)
        
        # 拷贝数据到 GPU
        cuda.memcpy_htod(d_input, h_input)
        
        # 执行推理
        context.execute_v2([int(d_input), int(d_output)])
        
        # 拷贝结果到 CPU
        cuda.memcpy_dtoh(h_output, d_output)
        
        return h_output


class ONNXRuntimeInference:
    """ONNX Runtime 推理"""
    
    def __init__(self, model_path: str, 
                 execution_provider: str = 'CUDAExecutionProvider'):
        """
        Args:
            model_path: ONNX 模型路径
            execution_provider: 'CUDAExecutionProvider' | 'CPUExecutionProvider' | 'TensorrtExecutionProvider'
        """
        import onnxruntime as ort
        
        self.session = ort.InferenceSession(
            model_path,
            providers=[execution_provider]
        )
        
        # 获取输入输出信息
        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]
        
        # 性能统计
        self.inference_times = []
        
    def inference(self, input_data: np.ndarray) -> dict:
        """推理
        
        Args:
            input_data: (B, C, H, W) 输入数据
            
        Returns:
            {
                'outputs': dict,
                'latency_ms': float
            }
        """
        import time
        
        start_time = time.time()
        
        outputs = self.session.run(
            self.output_names,
            {self.input_name: input_data}
        )
        
        latency_ms = (time.time() - start_time) * 1000
        self.inference_times.append(latency_ms)
        
        return {
            'outputs': dict(zip(self.output_names, outputs)),
            'latency_ms': latency_ms
        }
    
    def get_stats(self) -> dict:
        """获取性能统计"""
        import numpy as np
        
        return {
            'avg_latency_ms': np.mean(self.inference_times),
            'max_latency_ms': np.max(self.inference_times),
            'min_latency_ms': np.min(self.inference_times),
            'p95_latency_ms': np.percentile(self.inference_times, 95),
            'fps': 1000 / np.mean(self.inference_times)
        }

Qualcomm NPU 部署

QNN (Qualcomm Neural Network) SDK

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
class QualcommNPUDeployer:
    """Qualcomm NPU 部署器"""
    
    def __init__(self, onnx_path: str):
        self.onnx_path = onnx_path
        
    def convert_to_dlc(self, output_path: str):
        """转换为 DLC 格式（Qualcomm 专用格式）"""
        # 使用 Qualcomm 的转换工具
        import subprocess
        
        cmd = [
            'snpe-onnx-to-dlc',
            '--input_network', self.onnx_path,
            '--output_path', output_path
        ]
        
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode != 0:
            print(f"Conversion failed: {result.stderr}")
            return None
        
        print(f"DLC model saved to: {output_path}")
        return output_path
    
    def quantize_for_npu(self, 
                         dlc_path: str,
                         calibration_data: list,
                         output_path: str):
        """为 NPU 量化"""
        import subprocess
        
        # 创建校准数据文件
        calib_file = 'calibration_list.txt'
        with open(calib_file, 'w') as f:
            for data in calibration_data:
                # 保存为 raw 文件
                raw_file = f"calib_{len(data)}.raw"
                data.astype(np.float32).tofile(raw_file)
                f.write(f"{raw_file}\n")
        
        cmd = [
            'snpe-dlc-quantize',
            '--input_dlc', dlc_path,
            '--input_list', calib_file,
            '--output_dlc', output_path
        ]
        
        result = subprocess.run(cmd, capture_output=True, text=True)
        
        if result.returncode != 0:
            print(f"Quantization failed: {result.stderr}")
            return None
        
        print(f"Quantized DLC saved to: {output_path}")
        return output_path

性能对比

不同平台推理速度

平台 模型 精度 延迟 FPS
NVIDIA Jetson Orin MobileNetV3 FP16 8 ms 125
NVIDIA Jetson Orin MobileNetV3 INT8 5 ms 200
Qualcomm QCS8255 MobileNetV3 INT8 12 ms 83
Raspberry Pi 5 MobileNetV3 INT8 55 ms 18
Google Coral Edge-TPU MobileNetV3 INT8 3 ms 333

模型大小对比

模型 FP32 FP16 INT8 压缩比
MobileNetV3 21 MB 10.5 MB 5.3 MB 4x
EfficientNet-B0 29 MB 14.5 MB 7.3 MB 4x
ResNet-18 45 MB 22.5 MB 11.3 MB 4x

最佳实践

优化优先级

优先级 优化技术 效果
P0 INT8 量化 4x 加速，4x 压缩
P1 剪枝 + 微调 2-3x 压缩
P2 知识蒸馏 保持精度
P3 模型架构搜索 架构优化

精度恢复技巧

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
# 量化感知训练（QAT）
class QuantizationAwareTrainer:
    """量化感知训练器"""
    
    def __init__(self, model, qconfig=None):
        from torch.quantization import get_default_qconfig, prepare_qat
        
        if qconfig is None:
            qconfig = get_default_qconfig('fbgemm')
        
        self.model = model
        self.model.qconfig = qconfig
        
        # 准备 QAT
        self.model = prepare_qat(model, inplace=True)
        
    def train_step(self, images, labels, optimizer):
        """训练一步"""
        self.model.train()
        
        optimizer.zero_grad()
        outputs = self.model(images)
        loss = F.cross_entropy(outputs, labels)
        loss.backward()
        optimizer.step()
        
        return loss.item()
    
    def convert_to_quantized(self):
        """转换为量化模型"""
        from torch.quantization import convert
        return convert(self.model.eval())

参考资源