TensorRT INT8量化实战：从模型转换到推理加速
最近在部署YOLOv5模型时遇到了性能瓶颈，决定尝试TensorRT的INT8量化来提升推理速度。以下是我的踩坑记录。
环境准备
pip install tensorrt
pip install torch torchvision
模型转换步骤
首先将PyTorch模型转换为ONNX格式：
import torch
model = torch.load('yolov5s.pt')
model.eval()
x = torch.randn(1, 3, 640, 640)
torch.onnx.export(model, x, 'yolov5s.onnx', export_params=True, opset_version=11)
然后使用TensorRT转换：
import tensorrt as trt
import os
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)
with open('yolov5s.onnx', 'rb') as f:
parser.parse(f.read())
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30
config.set_flag(trt.BuilderFlag.INT8)
# 设置量化校准
config.set_calibration_profile(builder.create_optimization_profile())
engine = builder.build_engine(network, config)
校准数据准备
import numpy as np
from torch.utils.data import DataLoader
# 准备100张图片用于校准
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)
calibration_data = []
for i, (images, _) in enumerate(val_loader):
if i >= 100:
break
calibration_data.append(images.numpy())
实际效果对比
在RTX 3080上测试：
- FP32: 45ms/次推理
- INT8: 22ms/次推理
- 加速比: 约2倍
踩坑总结
- 校准数据量不足：最初只用了10张图片，导致精度下降2%
- 内存溢出：未设置最大workspace，需要手动调整到1GB
- 精度损失：INT8量化后mAP从78%降到76%，但可接受
建议使用TensorRT 8.5+版本，兼容性更好。
讨论