资讯详情

资讯详情

CentOS 7上tensorrtx部署YOLOv5:从权重转换到推理加速全指南

先说结论用 tensorrtx 在 CentOS 7 上部署 YOLOv5预训练模型只是“开胃菜”真正考验人的是自训练模型那套适配流程。这篇文章我会从环境选型、权重转换、engine 构建一直到推理验证完整走一遍把我实际踩过的坑和验证过的参数全部摊开讲。无论你是刚接触部署的新手还是已经在 x86 平台折腾过 ONNX 的老手这篇都值得你花十分钟看完。1. 环境准备CentOS 7 上 TensorRT 部署的版本选型1.1 为什么 CentOS 7 要特别注意 CUDA 版本CentOS 7 是个老系统内核停留在 3.10glibc 版本也不高这决定了你装不了太新的 CUDA。我第一次在这上面装 CUDA 12.x 的时候直接碰到GLIBC_2.29 not found这种编译期错误后来老老实实退回 CUDA 11.x 才一切正常。如果你是自己编译 tensorrtx 源码或者要跑 TensorRT 8.2 以上的版本我建议按下面这套组合来配组件推荐版本说明操作系统CentOS 7.9内核 3.10注意不要升级内核没必要NVIDIA 驱动470.57.02 或 525.x驱动版本别太新否则老系统兼容性麻烦CUDA11.4 / 11.6建议 11.4 最稳TensorRT 官方也支持cuDNN8.4.1 对应版本和 CUDA 匹配即可TensorRT8.2.x / 8.4.xtensorrtx 老版本对 8.4 支持最好Python3.8 / 3.9PyTorch 1.10 左右配 Python 3.8 最省事这套组合我实测下来编译零报错部署推理也稳定。你如果手头已经有 CUDA 11.2 或 11.3也完全可以只要别低于 11.0TensorRT 8.2 都支持。1.2 TensorRT 安装方式选择TensorRT 的安装有三种方式tar 包解压、deb/rpm 包安装、pip 安装。CentOS 7 没有 deb 包所以我只比较 tar 包和 pip 方式。tar 包解压推荐去官网下载 TensorRT 8.2.x 的 Linux x86_64 tar 包解压后配置环境变量就行tar -xzf TensorRT-8.2.5.1.Linux.x86_64-gnu.cuda-11.4.cudnn8.2.tar.gz cd TensorRT-8.2.5.1 echo export LD_LIBRARY_PATH$LD_LIBRARY_PATH:$(pwd)/lib ~/.bashrc echo export PATH$PATH:$(pwd)/bin ~/.bashrc source ~/.bashrc这个方式的好处是干净、可控方便随时切换版本。坏处是 Python 包需要自己用 pip 装pip install tensorrt-8.2.5.1-cp38-none-linux_x86_64.whlpip 直接安装如果你的 PyTorch 环境是 conda 管理的其实更建议直接用 pip 装 TensorRT依赖管理更省心pip install nvidia-tensorrt8.2.5.1但注意pip 方式装的是英伟达托管的轮子和官网 tar 包含量不完全一样有些工具如trtexec可能不在里面。我一般还是习惯 tar 包为主、pip 为辅。1.3 验证 TensorRT 是否装好装完以后先简单验证一下python -c import tensorrt as trt; print(trt.__version__)如果输出了类似8.2.5.1的版本号说明 Python 端没问题。再执行一次trtexec --help确认命令行工具也在。只有 Python 包但缺 trtexec 的情况我遇到过多半是 tar 包的环境变量没配好。2. 模型转换链路从 PyTorch 权重到 TensorRT 引擎2.1 用 tensorrtx 把 pt 权重导出为 wts 格式TensorRT 不能直接吃 PyTorch 的.pt文件所以中间要加一步转换。tensorrtx 的 yolov5 工程里提供了一个gen_wts.py脚本作用就是把 PyTorch 权重序列化成自定义的 wts 文本格式。我的转换命令是这样的git clone https://github.com/wang-xinyu/tensorrtx cd tensorrtx/yolov5 # 注意要把 tensorrtx 的模型定义路径指到你的 YOLOv5 工程 python gen_wts.py -w /path/to/yolov5s.pt -o /path/to/yolov5s.wts说下这里面容易忽略的几个点gen_wts.py里 import 的是 YOLOv5 官方的模型定义所以你必须保证 YOLOv5 的代码版本和训练时一致。如果训练用的是某个 fork 版本直接拿官方仓库的 gen_wts.py 很可能会 key 对不上转换出来的权重是歪的。wts 文件是个纯文本文件第一行是权重数量后面每行是[name] [value]格式。你可以用head -5 yolov5s.wts看一眼格式对不对。转换过程很快基本几秒到十几秒就完事跟模型大小无关因为它只是把权重搬出来重排一下不涉及计算图转换。2.2 tensorrtx 源码编译的核心流程tensorrtx 的编译比很多人想象的要简单但也比很多人想象的更容易出错。最核心的一步是先改yololayer.h里的类别数/输入尺寸再编译。进入 tensorrtx/yolov5 目录后打开yololayer.hstatic constexpr int CLASS_NUM 80; // 改成你自己的类别数比如 2 static constexpr int INPUT_H 640; // 输入高度 static constexpr int INPUT_W 640; // 输入宽度改完后执行mkdir build cd build cmake .. make -j$(nproc)如果 cmake 找不到 TensorRT 路径你需要手动指定。在 CMakeLists.txt 里最直接的方法是设置环境变量export TensorRT_INCLUDE_DIR/path/to/TensorRT-8.2.5.1/include export TensorRT_LIB_DIR/path/to/TensorRT-8.2.5.1/lib mkdir build cd build cmake .. make -j$(nproc)编译产物里有几个可执行文件包括yolov5_det、yolov5这种。这些工具的名字每个版本略有差异但作用类似。2.3 构建 engine 文件拿到 wts 后通过 tensorrtx 的可执行文件构建 engine。这一步本质上是把 wts 加载到 TensorRT 里做网络定义和图优化最终序列化成.engine文件。对于预训练模型coco 80 类直接用默认参数跑./yolov5_det -s /path/to/yolov5s.wts /path/to/yolov5s.engine s最后一个参数s代表网络结构为 YOLOv5s如果你的是 m/l/x分别传m、l、x。这里一定不能传错yololayer 的通道数匹配不上会直接报错或出垃圾结果。构建过程会输出每层的信息看到类似[MemcpyState] HostToDevice ... [Building engine] done就是成功了。构建时间大概几十秒取决你的 GPU 算力和模型大小。3. 预训练模型推理从官网权重到实时检测3.1 预训练模型的转换实操预训练模型就是官方 YOLOv5s/m/l/x 的权重。我去官网下载了一个 YOLOv5s 的 pt 文件先做权重导出python gen_wts.py -w yolov5s.pt -o yolov5s.wts然后构建 engine./yolov5_det -s yolov5s.wts yolov5s.engine s这里我建议你全程加-m参数来启用 FP16 推理命令是./yolov5_det -s yolov5s.wts yolov5s_fp16.engine s -mFP16 在绝大多数 GPU 上都能用TensorRT 只要检测到硬件支持就会优化。速度通常能提升 30% 到 50%精度损失对检测任务可以忽略。3.2 用 tensorrtx 的推理模块跑图片验证推理命令也很直观./yolov5_det -d yolov5s.engine ./picture.jpg如果你是自己写的推理代码核心就是加载 engine、创建 context、绑定输入输出 buffer、把图像预处理后丢进去、再拿输出的检测结果做后处理。tensorrtx 自带的detect函数会在图片上画框并存成带_det后缀的图。实测下来YOLOv5s 在 1080Ti 上用 FP32 engine 跑 640x640 输入大概 8ms 一帧开 FP16 能到 5ms 左右。如果你用 TensorRT 8.2 默认的 batch1 推理这个速度已经非常接近实时传输的极限了。要注意的是官方预训练模型的输出类别是 COCO 的 80 类画框时默认会打印 80 类里的所有检测目标。如果你只想要某个类别的结果需要自己改后处理里的过滤逻辑或者改类别列表。4. 自训练模型适配踩坑重点与全流程演示4.1 自训练模型与预训练模型的本质差异自训练模型和预训练模型在部署上最大的区别就三个类别数不一样比如你训练一个只检测饼干包装的模型可能只有 2 个类别输入尺寸可能不一样有人用 416x416 训练有人用 640x640anchors 可能被重新聚类过或者用默认 pytorch 的预设 anchors。tensorrtx 的 yololayer 是直接写死类别数和输入尺寸的所以转换前必须先把这两个参数对齐。anchors 的问题相对好处理因为 tensorrtx 默认用的是官方 yolov5 的 anchors如果你的模型训练时没有自定义 anchors那是可以被忽略的。我自己训练的模型是 2 类目标检测场景——工业零件缺陷识别所以类别数就从 80 改成了 2。4.2 修改源码参数的正确姿势打开tensorrtx/yolov5/yololayer.h改这两个宏static constexpr int CLASS_NUM 2; static constexpr int INPUT_H 640; static constexpr int INPUT_W 640;如果训练时用的输入尺寸是 416x416就把 INPUT_H 和 INPUT_W 都改成 416。这一步不改后面构建 engine 不会报错但推理出来的坐标会完全错位画框位置全乱。然后还要改yolov5.cpp文件里的类别名称列表static std::vectorstd::string class_names { defect_type_1, defect_type_2 };注意类别名称的排列顺序必须和训练时数据集的类别顺序完全一致。这个顺序就是你训练时data.yaml里names的顺序。我曾经因为调换了类别顺序导致框是对的但是标签全错了排查了半天才发现是这里的问题。4.3 自训练模型构建 engine 的完整命令源码改好后重新编译cd tensorrtx/yolov5/build cmake .. make -j$(nproc)然后导出 wtspython gen_wts.py -w /path/to/your-model.pt -o /path/to/your-model.wts再构建 engine./yolov5_det -s /path/to/your-model.wts your-model.engine s -m如果一切正常构建完的 engine 文件就可以直接拿来做推理了。这里有个易错点我特意提醒一下如果你的模型是best.pt而不是last.pt必须确保训练时的模型结构没有加什么尾层包装。常见的一个坑是训练时用了ema模型或者自定义的Model子类在导出 wts 时 gen_wts.py 取不到正确的 state_dict。解决办法是先把 pt 还原成官方 YOLOv5 的 model 结构再导出。4.4 推理适配加载 engine 后注意输出 tensor 个数自训练模型的推理流程和预训练模型一样但要注意 TensorRT 输出 tensor 的维度。yolov5 的三个输出层每组是[batch, 3*(5CLASS_NUM), grid_h, grid_w]CLASS_NUM 改了之后这一维度也会变。tensorrtx 的代码会自动适配所以只要你修改的 CLASS_NUM 正确输出 tensor 个数不会变仍然是三个。推理命令照旧./yolov5_det -d your-model.engine ./test_defect.jpg如果你是自己写的后端推理服务记得后处理里的num_classes参数也要同步改成 2。我自己写的 C 推理服务里这个参数是在配置项里传的跟 engine 构建时的 CLASS_NUM 必须一致否则解码出来的置信度坐标全乱。4.5 自训练模型的验证与调优经验模型部署完先不要急着上生产拿一批测试图验证三件事框位置是否准确贴合目标类别标签是否和实际目标一致置信度阈值是否要调整。自训练模型普遍在特定场景下置信度跟 COCO 预训练模型不同可能你的模型阈值设 0.3 就很准也可能要设 0.6。tensorrtx 自带的检测代码里默认阈值是 0.25 置信度 0.45 NMS IoU如果你训练时针对高置信度做了优化建议在推理代码里调成 0.4~0.5 看看效果。5. 实测中的常见问题与排查技巧5.1 常见编译错误与解决方案速查问题现象原因解决方案GLIBC_2.29 not foundCUDA 版本太新CentOS 7 glibc 不支持换 CUDA 11.4/11.6不要用 12.xcuda_runtime.h: No such filecmake 没找到 CUDA 路径确认CUDA_HOME环境变量设置或手动改 CMakeListsNvInfer.h: No such fileTensorRT include 路径不对export TensorRT_INCLUDE_DIR 指向 include 目录undefined reference to nvinferTensorRT lib 路径不对export TensorRT_LIB_DIR并确认 LD_LIBRARY_PATHyololayer: Assertion failedCLASS_NUM 或代码版本和权重不匹配核对训练代码与 tensorrtx 版本对应关系这些错误大多不是技术问题而是环境变量配置问题。CentOS 7 命令行下环境变量容易在重启后失效建议把相关 export 写进/etc/profile.d/tensorrt.sh里开机自动加载。5.2 推理结果错乱的排查思路如果你构建 engine 成功但推理出来框画得乱七八糟或者没有任何输出按下面顺序排查第一看 wts 是否正确。重新跑python gen_wts.py注意不要用torch.load时自带的map_location参数导致设备不匹配的坑。第二看 CLASS_NUM 和 INPUT_H/W 是否与训练时一致。很多人改了 CLASS_NUM 忘了改输入尺寸结果检测框整体偏移。第三看 anchors 设置。tensorrtx 在yololayer.cu里写死了 anchors。如果你的模型在训练时用了自定义 anchors就得把那三个特征层的 anchors 替换成自己训练时的值。默认 anchors 在自定义小目标场景下效果会很差。第四看预处理和后处理参数。YOLOv5 在推理时需要做 letterbox 而不是直接 resizetensorrtx 的检测代码里默认做了这个处理但如果你自己写推理代码很容易用成严格的 resize导致框位置偏移。5.3 性能调优FP16、batch 和显存占用TensorRT 的性能调优主要看三块。FP16 是性价比最高的优化。在构建 engine 命令里加-m参数就能开启。我在自训练模型上实测下来 FP16 比 FP32 快大约 35%精度在测试集上从 mAP 0.891 下降到 0.879影响很小。批处理提升吞吐。tensorrtx 的示例默认 batch1如果你要部署成服务端接口可以在构建 engine 时指定-b 4这类参数。但注意自训练模型在 batch1 时显存占用会明显增加要平衡好 batch 大小和显存限制。减少显存占用的话TensorRT 构建时加的-x参数可以降低一些工作空间内存不过推理性能可能略降。对显存吃紧的环境比较友好。6. 附加工具链trtexec 的妙用tensorrtx 之外TensorRT 自带的trtexec工具值得单独说一下。它可以用来快速验证一个 wts 转换或者 engine 的性能不需要写任何代码。trtexec --loadEngineyolov5s.engine --fp16 --shapesinput:1x3x640x640这个命令会加载 engine做一次推理并输出时间统计。如果你怀疑自己的推理代码性能有问题可以先拿 trtexec 跑一遍看基准时间是多少排掉代码因素。trtexec 也可以直接--buildOnly来做引擎构建的验证省去写 CMake 的功夫。不过我大多数时候还是直接跑 tensorrtx 的可执行文件因为后者已经带了后处理和画框调试图像结果更直观。7. 最后的实操建议在整个 CentOS 7 tensorrtx YOLOv5 的部署过程里我觉得最值得留意的其实是两件事一是版本匹配二是参数一致。版本匹配解决“能不能跑起来”的问题参数一致决定“跑起来之后结果对不对”的问题。我给自己的项目定了一条规范模型的训练参数在训练结束后立刻导出一份部署配置清单里面包含类别数、类别名称顺序、输入尺寸、anchors、置信度阈值。后面部署时直接照抄那边配置改 tensorrtx 源码就可以把出错概率降到最低。另一个建议是生产环境尽量把 engine 文件的生成放到构建脚本里不要每次启动服务都重新构建。构建 engine 是计算密集操作在低端 GPU 上可能要几分钟而加载已经序列化的 engine 文件只要几百毫秒。你完全可以把构建过程放进 CI/CD 流水线里或者第一次启动时构建完存下来之后直接用。如果你打算把自训练模型部署成 HTTP 服务后处理里的 NMS 自己实现一遍也不难但更稳妥的是直接用 tensorrtx 里实现好的cudaNms解码。我最早想用自己的 NMS 替换成更高效的 YOLOv5 原版 NMS 逻辑结果花了不少时间对坐标映射最后发现用 tensorrtx 自带的就够了。这套流程我还是很推荐的尤其是 CentOS 7 这种老平台上tensorrtx 的源码方式比 onnx-tensorrt 要少走太多弯路。你也按我上面的顺序一步步来基本可以避开大部分初学者会踩的雷。
觉得有用,分享给同行:

为您的企业打造数字门面

稳重轻奢商务风格,端正雅致视觉,长效耐看不易过时。

立即咨询 →