好的,阿杰,我们来看 TensorRT C++ 部署 的完整思路、步骤和注意事项。我会把流程拆解清楚,适合直接在工程中实践。


1️⃣ 前提条件

  1. 安装 NVIDIA 驱动(对应你的 GPU 型号和 CUDA 版本)
  2. 安装 CUDA Toolkit
  3. 安装 cuDNN
  4. 安装 TensorRT SDK(包含 libnvinferlibnvinfer_plugininclude 等)
  5. C++ 编译器支持(Linux 推荐 GCC ≥ 7,Windows 推荐 MSVC 2019+)

参考官方安装指南:TensorRT Installation Guide


2️⃣ 基本部署流程

2.1 导出模型

  • TensorRT 支持 ONNX、Caffe 等模型格式
  • 对 PyTorch 模型: import torch dummy_input = torch.randn(1, 3, 224, 224).cuda() torch.onnx.export(model, dummy_input, "model.onnx", opset_version=13)
  • 对 TensorFlow 模型:可先转换为 SavedModel 再用 tf2onnx 转 ONNX

参考:ONNX Export from PyTorch


2.2 使用 TensorRT 生成 engine 文件

示例 C++ 代码:

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

using namespace nvinfer1;

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

    // 创建 Builder
    IBuilder* builder = createInferBuilder(logger);
    INetworkDefinition* network = builder->createNetworkV2(0U);

    auto parser = nvonnxparser::createParser(*network, logger);
    if (!parser->parseFromFile("model.onnx", static_cast<int>(ILogger::Severity::kWARNING))) {
        std::cerr << "ONNX parse failed!" << std::endl;
        return -1;
    }

    // Builder 配置
    IBuilderConfig* config = builder->createBuilderConfig();
    config->setMaxWorkspaceSize(1 << 30); // 1GB
    config->setFlag(BuilderFlag::kFP16);  // 开启 FP16 加速

    // 生成 engine
    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
    if (!engine) {
        std::cerr << "Build engine failed!" << std::endl;
        return -1;
    }

    // 保存 engine
    IHostMemory* modelStream = engine->serialize();
    std::ofstream file("model.trt", std::ios::binary);
    file.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());

    // 释放资源
    modelStream->destroy();
    engine->destroy();
    network->destroy();
    builder->destroy();
    return 0;
}

参考官方文档:TensorRT C++ API


2.3 推理部署

加载 engine 并执行推理:

#include <NvInfer.h>
#include <cuda_runtime_api.h>
#include <vector>

int main() {
    // 1. 反序列化 engine
    std::ifstream file("model.trt", std::ios::binary);
    file.seekg(0, std::ios::end);
    size_t size = file.tellg();
    file.seekg(0, std::ios::beg);
    std::vector<char> engineData(size);
    file.read(engineData.data(), size);

    nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), size);

    // 2. 创建执行上下文
    nvinfer1::IExecutionContext* context = engine->createExecutionContext();

    // 3. 准备输入输出缓冲
    void* buffers[2];
    cudaMalloc(&buffers[0], inputSize);
    cudaMalloc(&buffers[1], outputSize);

    // 4. 将输入数据拷贝到 GPU
    cudaMemcpy(buffers[0], hostInput, inputSize, cudaMemcpyHostToDevice);

    // 5. 推理
    context->executeV2(buffers);

    // 6. 拷贝输出回 CPU
    cudaMemcpy(hostOutput, buffers[1], outputSize, cudaMemcpyDeviceToHost);

    // 7. 释放资源
    cudaFree(buffers[0]);
    cudaFree(buffers[1]);
    context->destroy();
    engine->destroy();
    runtime->destroy();

    return 0;
}

参考:TensorRT Inference Example


3️⃣ 部署优化建议

  1. 使用 FP16/INT8:减少显存和计算时间
  2. 批量推理:使用 batch,提高 GPU 利用率
  3. 使用 CUDA Stream:异步推理,提高吞吐
  4. TensorRT Plugin:自定义操作支持非标准层
  5. 序列化 engine:避免每次启动重新 build,提高启动速度

4️⃣ 参考资料

  1. TensorRT Installation Guide
  2. TensorRT C++ API
  3. TensorRT Developer Guide
  4. ONNX Export from PyTorch