YOLOv8 实时推理库：TensorRT 8 到 TensorRT 10 适配指南

发表评论

概述

本次更新为 YOLOv8 实时推理库添加了对 CUDA 12.9 和 TensorRT 10 的完整支持，同时兼容原有的 CUDA 11.8 + TensorRT 8 架构。通过 NV_TENSORRT_MAJOR 宏定义（NV_TENSORRT_MAJOR == 8 为 TensorRT 8，NV_TENSORRT_MAJOR == 10 为 TensorRT 10）在源代码中条件编译不同版本的 API 调用。原TensorRT 8的代码来源于https://github.com/wang-xinyu/tensorrtx.git

核心适配策略

条件编译宏设计

在 model.cpp 等文件中使用 #if NV_TENSORRT_MAJOR 指令统一管理两个版本的差异：

#if NV_TENSORRT_MAJOR < 10
    // TensorRT 8 代码路径
    builder->createNetworkV2(0U);
#else
    // TensorRT 10 代码路径
    uint32_t flags = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    builder->createNetworkV2(flags);
#endif

通过 NV_TENSORRT_MAJOR 宏在编译时自动选择对应版本的 API，无需运行时判断。

主要适配内容

1. 网络定义创建方式

TensorRT 8（旧版本）：

nvinfer1::INetworkDefinition* network = builder->createNetworkV2(0U);

TensorRT 10（新版本）：

uint32_t flags = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(flags);

变更原因：TensorRT 10 引入了显式批处理（Explicit Batch）要求，必须在创建网络时设置 kEXPLICIT_BATCH 标志。

2. 输入张量维度定义

TensorRT 8：

nvinfer1::ITensor* data = network->addInput(kInputTensorName, dt,
    nvinfer1::Dims3{3, kInputH, kInputW});  // CHW 格式

TensorRT 10：

nvinfer1::ITensor* data = network->addInput(kInputTensorName, dt,
    nvinfer1::Dims{4, {kBatchSize, 3, kInputH, kInputW}});  // NCHW 格式

变更原因：TensorRT 10 强制使用 NCHW（N为batch维度）格式，而 TensorRT 8 支持隐式批处理的 CHW 格式。

3. 模型构建接口变更

3.1 调整大小层（Resize Layer）

TensorRT 8：

upsample10->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
upsample10->setScales(scale, 3);  // scale 数组长度为3

TensorRT 10：

upsample10->setResizeMode(nvinfer1::InterpolationMode::kNEAREST);
upsample10->setScales(scale, 4);  // scale 数组长度为4，新增 batch 维度

变更原因：TensorRT 10 的 resize 层需为 batch 维度提供插值系数，scale 数组从 [1.0, 2.0, 2.0] 扩展为 [1.0, 1.0, 2.0, 2.0]。

3.2 卷积层（Convolution Layer）

TensorRT 8：

conv22_cv3_0_2->setStride(nvinfer1::DimsHW{1, 1});
conv22_cv3_0_2->setPadding(nvinfer1::DimsHW{0, 0});

TensorRT 10：

conv22_cv3_0_2->setStrideNd(nvinfer1::Dims{2, {1, 1}});
conv22_cv3_0_2->setPaddingNd(nvinfer1::Dims{2, {0, 0}});

变更原因：

  • setStride/Padding → setStrideNd/PaddingNd：多维接口更灵活
  • 参数从 DimsHW 改为 Dims，支持任意维度组合

3.3 全连接层（Fully Connected Layer）→ 卷积层

TensorRT 8：

nvinfer1::IFullyConnectedLayer* yolo = network->addFullyConnected(
    *pool2->getOutput(0), kClsNumClass,
    weightMap["model.9.linear.weight"], weightMap["model.9.linear.bias"]);

TensorRT 10：

nvinfer1::IConvolutionLayer* yolo = network->addConvolutionNd(
    *pool2->getOutput(0),
    kClsNumClass,
    nvinfer1::DimsHW{1, 1},
    weightMap["model.9.linear.weight"], weightMap["model.9.linear.bias"]);

变更原因：TensorRT 10 废弃了 addFullyConnected，统一使用 addConvolutionNd 实现。

3.4 池化层输出维度

TensorRT 8：

nvinfer1::IPoolingLayer* pool2 = network->addPoolingNd(
    *conv_class->getOutput(0), nvinfer1::PoolingType::kAVERAGE,
    nvinfer1::DimsHW{dims.d[1], dims.d[2]});

TensorRT 10：

nvinfer1::IPoolingLayer* pool2 = network->addPoolingNd(
    *conv_class->getOutput(0), nvinfer1::PoolingType::kAVERAGE,
    nvinfer1::DimsHW{dims.d[2], dims.d[3]});

变更原因：维度索引适配 NCHW 格式。

4. DFL（Distribution Focal Loss）层适配

4.1 Reshape 维度

TensorRT 8：

shuffle1->setReshapeDimensions(nvinfer1::Dims3{4, 16, grid});
shuffle1->setSecondTranspose(nvinfer1::Permutation{1, 0, 2});

TensorRT 10：

shuffle1->setReshapeDimensions(nvinfer1::Dims4{0, 4, 16, grid});
shuffle1->setSecondTranspose(nvinfer1::Permutation{0, 2, 1, 3});

变更原因：

  • Dims 从 3D 扩展为 4D（添加 batch 维度）
  • Permutation 索引适配 4D 张量

4.2 SoftMax 轴设置

TensorRT 10 新增：

softmax->setAxes(1U << 1);  // 在第2维（channel维度）计算 SoftMax

变更原因：TensorRT 10 引入了 setAxes 接口来指定 SoftMax 的计算轴，而 TensorRT 8 在创建层时隐式确定。

4.3 后续 Shuffle Reshape

TensorRT 8：

shuffle2->setReshapeDimensions(nvinfer1::Dims2{4, grid});

TensorRT 10：

shuffle2->setReshapeDimensions(nvinfer1::Dims3{0, 4, grid});

变更原因：添加 batch 维度。

5. Concatenation 层轴设置

TensorRT 10 新增：

cat22_dfl_0->setAxis(1);
cat22_dfl_1->setAxis(1);
cat22_dfl_2->setAxis(1);

变更原因：TensorRT 10 需要显式指定拼接的维度，默认为轴 1（channel 维度）。

6. Shuffl e 层和 Slice 层

TensorRT 8：

shuffle22_0->setReshapeDimensions(nvinfer1::Dims2{64 + kNumClass,
    (kInputH / strides[0]) * (kInputW / strides[0])});

nvinfer1::ISliceLayer* split22_0_0 = network->addSlice(
    *shuffle22_0->getOutput(0), nvinfer1::Dims2{0, 0},
    nvinfer1::Dims2{64, (kInputH / strides[0]) * (kInputW / strides[0])},
    nvinfer1::Dims2{1, 1});

TensorRT 10：

shuffle22_0->setReshapeDimensions(nvinfer1::Dims3{kBatchSize, 64 + kNumClass,
    (kInputH / strides[0]) * (kInputW / strides[0])});

nvinfer1::ISliceLayer* split22_0_0 = network->addSlice(
    *shuffle22_0->getOutput(0), nvinfer1::Dims3{0, 0, 0},
    nvinfer1::Dims3{kBatchSize, 64, (kInputH / strides[0]) * (kInputW / strides[0])},
    nvinfer1::Dims3{1, 1, 1});

变更原因：所有维度添加 batch 维度（维度索引从 1D → 2D，2D → 3D）。

7. Builder 配置接口

TensorRT 8：

builder->setMaxBatchSize(kBatchSize);
config->setMaxWorkspaceSize(16 * (1 << 20));

TensorRT 10：

config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 16 * (1 << 20));

变更原因：

  • setMaxBatchSize 移至 IBuilder 接口
  • setMaxWorkspaceSize 改为 setMemoryPoolLimit，类型为 MemoryPoolType::kWORKSPACE

8. 推理执行接口

8.1 反序列化接口更新

TensorRT 8：

deserialize_engine(std::wstring &engine_name,
    nvinfer1::IRuntime** runtime,
    nvinfer1::ICudaEngine** engine,
    nvinfer1::IExecutionContext** context);

TensorRT 10：

deserialize_engine(std::wstring &engine_name,
    nvinfer1::IRuntime** runtime,
    nvinfer1::ICudaEngine** engine,
    nvinfer1::IExecutionContext** context,
    nvinfer1::ICudaEngine** engine_10);  // 新增参数

变更原因：需要分离 TensorRT 8 和 10 的引擎接口（虽然尚未实际使用）。

8.2 输出张量获取方式

TensorRT 8：

auto out_dims = engine->getBindingDimensions(1);
model_bboxes = out_dims.d[0];

TensorRT 10：

const char* output_name = engine->getIOTensorName(1);
auto out_dims = engine->getTensorShape(output_name);
model_bboxes = out_dims.d[0];

变更原因：

  • TensorRT 10 不再支持通过索引绑定获取维度
  • 必须通过 getIOTensorName() 获取张量名称

8.3 内存缓冲区绑定方式

TensorRT 8：

context.enqueue(batchsize, buffers, stream, nullptr);

TensorRT 10：

context.setInputShape(kInputTensorName, nvinfer1::Dims{4, {batchsize, 3, kInputH, kInputW}});
context.setTensorAddress(kInputTensorName, buffers[0]);
context.setTensorAddress(kOutputTensorName, buffers[1]);
context.enqueueV3(stream);

变更原因：

  • 推理方法从 enqueue 改为 enqueueV3（异步执行接口）
  • 需要显式设置输入输出张量的地址和形状
  • setInputShape 在 enqueueV3 之前调用

8.4 绑定数量断言

TensorRT 8：

assert(engine->getNbBindings() == 2);
const int inputIndex = engine->getBindingIndex(kInputTensorName);
const int outputIndex = engine->getBindingIndex(kOutputTensorName);

TensorRT 10：

// 删除 getNbBindings() 检查
const char* input_name = engine->getIOTensorName(0);
const char* output_name = engine->getIOTensorName(1);

变更原因：TensorRT 10 不再使用 “bindings” 概念，改用张量（tensor）概念。

兼容性文件结构

头文件 types.h、yolov8det.h、yolov8cls.h、yolov8seg.h 保持不变，确保上层接口无需修改。

条件编译宏 NV_TENSORRT_MAJOR 通过编译器宏定义传入：

  • TensorRT 8：NV_TENSORRT_MAJOR 定义为 8
  • TensorRT 10：NV_TENSORRT_MAJOR 定义为 10

总结

本次适配通过 #if NV_TENSORRT_MAJOR 条件编译策略，实现了从 TensorRT 8 到 TensorRT 10 的平滑迁移：

✅ 显式批处理（Explicit Batch）支持
✅ NCHW 输入格式统一
✅ API 接口现代化（enqueueV3、setTensorAddress）
✅ 向后兼容旧版本（通过宏切换）
✅ 清晰的代码分支管理

开发人员只需在编译时定义 NV_TENSORRT_MAJOR 为 10，即可自动使用 TensorRT 10 的所有 API，无需手动修改业务逻辑。

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注