模型优化与部署

1. 模型优化概述

模型优化旨在在保持精度的前提下，减少模型大小、降低延迟、减少功耗。

精度与效率是模型优化的两端，量化、剪枝、蒸馏等技术在两者之间寻找平衡：压缩模型体积、加速推理速度，以尽可能小的精度损失换取效率提升。

2. 量化

2.1 量化原理

将浮点参数映射到低精度整数：

$x_q = \text{round}\left(\frac{x}{S}\right) + Z$

其中 $S$ 为缩放因子（Scale），$Z$ 为零点（Zero Point）。

反量化：

$x \approx S \cdot (x_q - Z)$

2.2 量化类型

2.3 PTQ（训练后量化）

import torch.quantization as quant

# 静态量化
model.eval()
model.qconfig = quant.get_default_qconfig('fbgemm')
prepared = quant.prepare(model)
# 用校准数据运行
with torch.no_grad():
    for batch in calib_loader:
        prepared(batch)
quantized = quant.convert(prepared)

2.4 QAT（量化感知训练）

model.train()
model.qconfig = quant.get_default_qat_qconfig('fbgemm')
prepared = quant.prepare_qat(model)

# 正常训练
for epoch in range(num_epochs):
    for batch in train_loader:
        loss = train_step(prepared, batch)

quantized = quant.convert(prepared)

伪量化：在前向传播中模拟量化效果，反向传播使用STE（Straight-Through Estimator）。

2.5 量化精度对比

3. 剪枝

3.1 剪枝类型

3.2 剪枝策略

幅度剪枝：移除绝对值最小的权重

$\text{mask}_{ij} = \begin{cases} 1 & |w_{ij}| > \tau \\ 0 & |w_{ij}| \leq \tau \end{cases}$

迭代剪枝：

1. 训练模型至收敛
2. 剪枝：移除p%最小权重
3. 微调：恢复精度
4. 重复2-3直到目标稀疏度

3.3 结构化剪枝

# 基于L1范数的通道剪枝
import torch.nn.utils.prune as prune

# 非结构化
prune.l1_unstructured(model.fc1, name='weight', amount=0.3)

# 结构化（整个通道）
prune.ln_structured(model.fc1, name='weight', amount=0.3, n=2, dim=0)

3.4 稀疏训练

RigL：周期性移除和生长权重

1. 前向+反向传播
2. 移除最小幅度的权重
3. 在梯度最大的位置生长新权重
4. 保持恒定稀疏度

4. 知识蒸馏

4.1 基本原理

用**大模型（教师）指导小模型（学生）**训练：

$\mathcal{L} = \alpha \cdot \mathcal{L}_{CE}(y, y_s) + (1-\alpha) \cdot T^2 \cdot \mathcal{L}_{KL}(\sigma(z_t/T), \sigma(z_s/T))$

4.2 温度Softmax

$\sigma(z_i/T) = \frac{e^{z_i/T}}{\sum_j e^{z_j/T}}$

• $T > 1$：软化概率分布，传递更多”暗知识”
• $T = 1$：标准Softmax

4.3 蒸馏变体

5. ONNX与部署

5.1 ONNX导出

import torch.onnx

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy_input, "model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch'}, 'output': {0: 'batch'}},
    opset_version=14
)

5.2 ONNX Runtime

import onnxruntime as ort

session = ort.InferenceSession("model.onnx")
input_name = session.get_inputs()[0].name
output = session.run(None, {input_name: input_data})

5.3 TensorRT

# ONNX → TensorRT
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)

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
engine = builder.build_serialized_network(network, config)

5.4 部署方案对比