C++与TensorRT实战：从零搭建高性能推理引擎的避坑指南

1次阅读

共计 4073 个字符，预计需要花费 11 分钟才能阅读完成。

在实际生产环境中，我们经常遇到 PyTorch 或 TensorFlow 等框架的 C ++ 推理性能不足的问题。特别是在实时视频分析、自动驾驶、工业质检等延迟敏感型场景中，原生框架的表现往往难以满足需求。

以一个典型的 ResNet50 图像分类任务为例，在相同硬件环境下：

PyTorch(C++): 约 45ms/ 帧
TensorRT(FP16): 约 8ms/ 帧
TensorRT(INT8): 约 5ms/ 帧

特性	TensorRT	ONNX Runtime	TVM
推理延迟	★★★★★	★★★☆	★★★☆
内存占用	★★★★☆	★★★★	★★★☆
部署复杂度	★★★☆	★★★★	★★☆
算子支持度	★★★★	★★★★☆	★★★☆
量化支持	★★★★★	★★★☆	★★★★

# PyTorch 转 ONNX
torch.onnx.export(model, dummy_input, "model.onnx", 
                 input_names=["input"], 
                 output_names=["output"],
                 dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})

# ONNX 转 TensorRT
trtexec --onnx=model.onnx --saveEngine=model.trt --fp16

// 创建 builder
auto builder = nvinfer1::createInferBuilder(logger);

// 创建 network
auto network = builder->createNetworkV2(1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));

// 创建 parser
auto parser = nvonnxparser::createParser(*network, logger);
parser->parseFromFile(onnxModelPath.c_str(), static_cast<int>(ILogger::Severity::kWARNING));

// 配置 builder
auto config = builder->createBuilderConfig();
config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 1 << 30); // 1GB
if (fp16) config->setFlag(BuilderFlag::kFP16);

// 创建校准器
class Int8EntropyCalibrator : public IInt8EntropyCalibrator2 {
public:
    Int8EntropyCalibrator(const std::vector<std::string>& imagePaths, int batchSize)
        : mImagePaths(imagePaths), mBatchSize(batchSize) {}

    int getBatchSize() const noexcept override { return mBatchSize;}

    bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override {
        // 实现数据加载逻辑
        return true;
    }
};

// 配置 INT8
config->setFlag(BuilderFlag::kINT8);
auto calibrator = std::make_unique<Int8EntropyCalibrator>(calibrationData, batchSize);
config->setInt8Calibrator(calibrator.get());

nsys profile --stats=true ./your_inference_app

输出示例：

Time(%)  Total Time  Calls  Avg       Min       Max  Name
50.2     120.45ms    100   1.2045ms  1.10ms    1.35ms  convolution
30.1     72.12ms     100   0.7212ms  0.70ms    0.75ms  matrixMultiply

// 手动选择 tactic
config->setTacticSources(1U << static_cast<uint32_t>(TacticSource::kCUBLAS));

// 禁用不需要的 precision
config->clearFlag(BuilderFlag::kTF32);

// 预分配输入输出 buffer
std::vector<void*> buffers(engine->getNbBindings());
for (int i = 0; i < engine->getNbBindings(); ++i) {auto dims = context->getBindingDimensions(i);
    size_t size = getSize(dims) * sizeof(float);
    cudaMalloc(&buffers[i], size);
}

不支持的算子
使用 plugin 机制实现自定义层
考虑用等效算子组合替换
回退到原生框架部分执行
版本兼容性
保持 CUDA/cuDNN/TensorRT 版本匹配
使用容器化部署 (Docker)
为不同版本维护多个引擎文件

动态 shape 处理

// 设置优化 profile
auto profile = builder->createOptimizationProfile();
profile->setDimensions(inputName, OptProfileSelector::kMIN, Dims4{1,3,224,224});
profile->setDimensions(inputName, OptProfileSelector::kOPT, Dims4{8,3,224,224});
profile->setDimensions(inputName, OptProfileSelector::kMAX, Dims4{32,3,224,224});
config->addOptimizationProfile(profile);

在医疗影像分析场景中，如何权衡 INT8 量化的精度损失与推理速度提升？
当遇到不支持动态 batch 的硬件平台时，有哪些替代方案可以实现类似效果？
在多模型流水线场景中，如何设计资源分配策略来最大化整体吞吐量？

├── CMakeLists.txt
├── include
│   ├── inference_engine.h
│   └── utils.h
├── src
│   ├── main.cpp
│   └── trt_inference.cpp
└── models
    ├── model.onnx
    └── model.trt

CMake 关键配置：

find_package(CUDA REQUIRED)
find_package(TensorRT REQUIRED)

add_executable(inference_app src/main.cpp src/trt_inference.cpp)
target_link_libraries(inference_app ${TensorRT_LIBRARY} ${CUDA_LIBRARIES})

模型	框架	精度	延迟 (ms)	显存 (MB)
ResNet50	PyTorch	FP32	45.2	1250
ResNet50	TensorRT	FP16	8.1	780
ResNet50	TensorRT	INT8	4.9	520
YOLOv5s	ONNX Runtime	FP32	32.7	980
YOLOv5s	TensorRT	INT8	6.3	450

注：测试环境为 T4 GPU，batch_size=8

启用详细日志：

class Logger : public nvinfer1::ILogger {void log(Severity severity, const char* msg) noexcept override {if (severity <= Severity::kWARNING)
            std::cout << msg << std::endl;
    }
} logger;

检查引擎有效性：

auto engine = runtime->deserializeCudaEngine(engineData.data(), engineSize);
if (!engine) {throw std::runtime_error("Failed to deserialize engine");
}

验证输入输出：

for (int i = 0; i < engine->getNbBindings(); ++i) {
    std::cout << "Binding" << i << ":" 
              << (engine->bindingIsInput(i) ? "Input" : "Output") 
              << ", Type:" << engine->getBindingDataType(i) 
              << ", Dims:" << engine->getBindingDimensions(i) 
              << std::endl;
}