C++使用TensorRT加速推理：从模型转换到性能优化的完整指南

1次阅读

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

深度学习模型在推理阶段常常面临性能瓶颈，尤其是在生产环境中需要处理大量请求时。原生框架如 PyTorch 或 TensorFlow 虽然提供了方便的模型训练接口，但在推理优化上往往不够高效。TensorRT 作为 NVIDIA 推出的高性能推理引擎，通过层融合、内核自动调优和量化等技术，可以显著提升推理速度。

TensorRT 的加速原理主要包括：

层融合（Layer Fusion）：将多个连续的操作合并为一个更高效的操作，减少内存访问和内核启动开销。
内核自动调优（Kernel Auto-Tuning）：根据目标硬件选择最优的内核实现。
量化（Quantization）：支持 FP16 和 INT8 精度，在保持模型精度的同时大幅提升计算效率。

在使用 TensorRT 之前，需要确保 CUDA、cuDNN 和 TensorRT 的版本匹配。以下是推荐的版本组合：

CUDA 11.x
cuDNN 8.x
TensorRT 8.x

安装步骤：

从 NVIDIA 官网下载对应版本的 CUDA Toolkit 并安装。
下载与 CUDA 版本匹配的 cuDNN 库，解压后复制到 CUDA 安装目录。

下载 TensorRT 的 tar 包，解压后设置环境变量：

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/tensorrt/lib

TensorRT 支持通过 ONNX 格式转换模型。以下是将 ONNX 模型转换为 TensorRT 引擎的步骤：

使用 PyTorch 或 TensorFlow 将训练好的模型导出为 ONNX 格式。

使用 TensorRT 的 onnx2trt 工具或 Python API 进行转换：

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)

with open("model.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))

config = builder.create_builder_config()
config.max_workspace_size = 1 << 30  # 1GB
engine = builder.build_engine(network, config)

以下是一个完整的 C ++ 推理示例，使用现代 C ++ 特性：

#include <NvInfer.h>
#include <NvOnnxParser.h>
#include <iostream>
#include <fstream>
#include <vector>

using namespace nvinfer1;

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

std::vector<char> loadEngine(const std::string& enginePath) {std::ifstream engineFile(enginePath, std::ios::binary);
    engineFile.seekg(0, std::ios::end);
    size_t size = engineFile.tellg();
    engineFile.seekg(0, std::ios::beg);
    std::vector<char> engineData(size);
    engineFile.read(engineData.data(), size);
    return engineData;
}

int main() {
    Logger logger;
    std::vector<char> engineData = loadEngine("model.engine");

    IRuntime* runtime = createInferRuntime(logger);
    ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), engineData.size());
    IExecutionContext* context = engine->createExecutionContext();

    // 分配输入输出显存
    void* buffers[2];
    int inputIndex = engine->getBindingIndex("input");
    int outputIndex = engine->getBindingIndex("output");

    cudaMalloc(&buffers[inputIndex], inputSize);
    cudaMalloc(&buffers[outputIndex], outputSize);

    // 执行推理
    context->executeV2(buffers);

    // 释放资源
    cudaFree(buffers[inputIndex]);
    cudaFree(buffers[outputIndex]);
    context->destroy();
    engine->destroy();
    runtime->destroy();
    return 0;
}

TensorRT 会自动尝试融合层，但可以通过配置进一步优化：

config->setFlag(BuilderFlag::kFP16);  // 启用 FP16 加速
config->setFlag(BuilderFlag::kREFIT);  // 允许后续重新优化

FP16 量化可以显著提升性能，而 INT8 量化则需要校准：

config->setFlag(BuilderFlag::kINT8);

// 校准器实现
class MyCalibrator : public IInt8EntropyCalibrator2 {
public:
    MyCalibrator() {}
    int getBatchSize() const noexcept override { return 1;}
    bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept override {
        // 提供校准数据
        return false;
    }
};

config->setInt8Calibrator(new MyCalibrator());

以下是一个简单的性能对比（单位：ms）：

Framework	Latency (ms)	Throughput (FPS)
PyTorch	50	20
TensorRT	10	100

模型转换失败：检查 ONNX 模型是否包含不支持的算子，或尝试更新 TensorRT 版本。

动态 shape 处理 ：使用profile 设置动态维度范围：

IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions("input", OptProfileSelector::kMIN, Dims4{1, 3, 224, 224});
profile->setDimensions("input", OptProfileSelector::kOPT, Dims4{4, 3, 224, 224});
profile->setDimensions("input", OptProfileSelector::kMAX, Dims4{8, 3, 224, 224});
config->addOptimizationProfile(profile);