1. 项目概述
这个项目展示了如何在Ubuntu系统上使用TensorRT的INT8量化技术,实现一个端到端的行人检测与人群密度分析系统。系统支持RTSP视频流输入,采用CMake构建工具管理整个工程。作为计算机视觉领域的典型应用场景,这种方案特别适合需要实时分析的智能监控、公共安全等边缘计算场景。
INT8量化是TensorRT的核心优化技术之一,能将模型体积压缩至原来的1/4,推理速度提升2-3倍,同时保持可接受的精度损失。我们选择行人检测和人群密度分析作为演示任务,因为这两个功能在实际应用中经常需要联合使用——先检测出每个人体,再统计区域密度。
需要模型API调用? 免费领10W Token,多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 环境准备与依赖安装
2.1 基础环境配置
推荐使用Ubuntu 20.04或22.04 LTS版本,这是目前最稳定的深度学习开发环境。需要预先安装:
bash复制sudo apt update && sudo apt install -y \
build-essential cmake git libopencv-dev \
libavcodec-dev libavformat-dev libswscale-dev \
libgstreamer-plugins-base1.0-dev libgstreamer1.0-dev
对于NVIDIA显卡驱动,建议使用官方.run文件安装最新版本。CUDA 11.x与TensorRT 8.x的兼容性最好,这是经过大量实践验证的稳定组合。
2.2 TensorRT安装要点
从NVIDIA官网下载对应CUDA版本的TensorRT本地.deb包。安装时特别注意:
bash复制sudo dpkg -i nv-tensorrt-repo-ubuntu2004-cuda11.4-trt8.2.5.1-ga-20220505_1-1_amd64.deb
sudo apt-key add /var/nv-tensorrt-repo-ubuntu2004-cuda11.4-trt8.2.5.1-ga-20220505/7fa2af80.pub
sudo apt update
sudo apt install tensorrt
验证安装是否成功:
bash复制dpkg -l | grep TensorRT
3. 模型转换与INT8量化
3.1 模型选择与转换
我们选用YOLOv5s作为行人检测模型,因其在精度和速度间取得了良好平衡。使用官方export.py脚本导出ONNX格式:
python复制python export.py --weights yolov5s.pt --include onnx --dynamic
然后通过TensorRT的trtexec工具转换为TensorRT引擎:
bash复制trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s.engine --fp16
3.2 INT8量化校准
INT8量化的核心是校准过程,需要准备约500张具有代表性的校准图像。创建校准器:
python复制class YOLOEntropyCalibrator(trt.IInt8EntropyCalibrator2):
def __init__(self, calibration_data):
self.cache_file = "yolov5s.calib"
self.data = calibration_data
self.current_index = 0
def get_batch(self, names):
if self.current_index + batch_size > len(self.data):
return None
batch = self.data[self.current_index:self.current_index+batch_size]
self.current_index += batch_size
return [np.ascontiguousarray(batch)]
量化转换命令:
bash复制trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s_int8.engine --int8 --calib=yolov5s.calib
4. RTSP视频流处理
4.1 视频流接入
使用FFmpeg处理RTSP流,解决常见的断流和延迟问题:
cpp复制AVFormatContext* fmt_ctx = NULL;
AVDictionary* options = NULL;
av_dict_set(&options, "rtsp_transport", "tcp", 0);
av_dict_set(&options, "stimeout", "5000000", 0); // 5秒超时
if (avformat_open_input(&fmt_ctx, rtsp_url.c_str(), NULL, &options) != 0) {
// 错误处理
}
4.2 视频解码优化
使用NVDEC硬件加速解码大幅提升性能:
cpp复制CUcontext cuContext;
cuCtxCreate(&cuContext, 0, device_id);
AVBufferRef* hw_device_ctx = av_hwdevice_ctx_alloc(AV_HWDEVICE_TYPE_CUDA);
AVHWDeviceContext* device_ctx = (AVHWDeviceContext*)hw_device_ctx->data;
AVCUDADeviceContext* cuda_ctx = device_ctx->hwctx;
cuda_ctx->cuda_ctx = cuContext;
5. CMake工程构建
5.1 项目结构设计
code复制├── CMakeLists.txt
├── include
│ ├── detector.h
│ ├── tracker.h
│ └── density_estimator.h
├── src
│ ├── main.cpp
│ ├── detector.cpp
│ └── density_estimator.cpp
└── third_party
├── tensorrt
└── opencv
5.2 关键CMake配置
cmake复制find_package(OpenCV REQUIRED)
find_package(TensorRT REQUIRED)
add_executable(pedestrian_analysis
src/main.cpp
src/detector.cpp
src/density_estimator.cpp)
target_include_directories(pedestrian_analysis PRIVATE
${OpenCV_INCLUDE_DIRS}
${TENSORRT_INCLUDE_DIR})
target_link_libraries(pedestrian_analysis
${OpenCV_LIBS}
nvinfer nvinfer_plugin cudart)
6. 性能优化技巧
6.1 内存管理
使用CUDA统一内存减少主机-设备间拷贝:
cpp复制void* unified_buffer;
cudaMallocManaged(&unified_buffer, buffer_size, cudaMemAttachGlobal);
6.2 流水线优化
设计三级处理流水线:
- 视频解码(NVDEC)
- 图像预处理(CUDA核函数)
- 模型推理(TensorRT)
使用CUDA流实现异步执行:
cpp复制cudaStream_t stream;
cudaStreamCreate(&stream);
context->enqueueV2(buffers, stream, nullptr);
7. 实际部署问题排查
7.1 常见问题与解决
| 问题现象 | 可能原因 | 解决方案 |
|---|---|---|
| 推理结果全零 | 输入数据范围错误 | 检查预处理是否匹配训练时的归一化 |
| 内存泄漏 | 未释放TRT资源 | 确保builder、runtime、engine正确释放 |
| RTSP断流 | 网络波动 | 增加重连机制和缓冲区 |
7.2 精度调优建议
INT8量化后若精度下降明显:
- 增加校准数据集多样性
- 尝试QAT(量化感知训练)
- 调整校准算法参数
8. 效果展示与性能指标
在RTX 3060显卡上的测试结果:
| 指标 | FP32 | FP16 | INT8 |
|---|---|---|---|
| 推理速度(FPS) | 45 | 78 | 112 |
| 内存占用(MB) | 1200 | 680 | 320 |
| mAP@0.5 | 0.72 | 0.71 | 0.68 |
人群密度分析算法采用基于检测的计数方法,在Mall数据集上达到92%的计数准确率。
这个项目的完整实现展示了从模型转换到实际部署的全流程,其中的技术要点也适用于其他计算机视觉任务的部署。在实际应用中,我们发现INT8量化虽然会带来约3-5%的精度下降,但在大多数监控场景中完全可接受,而获得的性能提升对于边缘设备至关重要。
