深度学习模型推理延迟优化踩坑记录
最近在优化一个BERT-base模型的推理延迟，从80ms优化到25ms，分享一下踩过的坑。
问题背景：原模型在GPU上推理耗时80ms，无法满足实时性要求。
第一轮优化：量化压缩
import torch
import torch.nn.quantized as nnq
# 量化模型
model = torch.quantization.quantize_dynamic(
model, {torch.nn.Linear}, dtype=torch.qint8
)
结果：延迟降到65ms，效果有限。
第二轮优化：TensorRT加速** 安装tensorrt后：
# 导出onnx
torch.onnx.export(model, input_tensor, "model.onnx")
# 构建engine
import tensorrt as trt
结果：延迟降到35ms，但推理精度下降1.2%。
第三轮优化：模型剪枝+蒸馏
from torch.nn.utils import prune
prune.l1_unstructured(model.linear, name='weight', amount=0.4)
# 蒸馏训练后精度恢复
结果：延迟25ms，精度保持在92%。
关键坑点：
- 量化前必须做充分的性能测试，避免精度损失过大
- TensorRT需要GPU支持，否则加速效果有限
- 剪枝后要进行蒸馏恢复，否则精度会大幅下降
完整测试环境：RTX 3090, CUDA 11.8, PyTorch 2.0
建议：先从模型剪枝开始，性价比最高。
讨论