📚 课程目录第2章 TensorRT部署环境搭建与基础推理

2. TensorRT部署环境搭建与基础推理

任务目标

在NVIDIA Jetson设备（本文以JetPack 5.1.2 / Ubuntu 20.04为例）上完成TensorRT的安装与验证，并运行一个基于TensorRT C++ API的最小推理程序。该程序读取一个ONNX格式的线性模型（y = 2x + 1），构建TensorRT引擎并执行推理，最终输出与ONNX Runtime参考结果一致的预测值。

环境与依赖

硬件：NVIDIA Jetson Orin / Xavier / Nano（本文适用于JetPack 5.x全系列）
系统：JetPack 5.1.2（含Ubuntu 20.04、CUDA 11.4、cuDNN 8.6）
TensorRT版本：8.5.2（随JetPack预装）
依赖库：libnvinfer-dev、libnvonnxparser-dev（JetPack默认已安装）
Python（用于生成测试模型）：3.8 + onnx 1.12 + onnxruntime 1.14

实现步骤

1. 验证TensorRT安装

JetPack自带TensorRT，无需单独下载。执行以下命令确认版本：

dpkg -l | grep nvinfer
# 预期输出包含 libnvinfer8 和 libnvinfer-dev，版本为 8.5.2.x

2. 生成测试用ONNX模型

在PC或Jetson上运行以下Python脚本，生成一个计算 y = 2x + 1 的线性模型。该模型仅用于验证TensorRT推理链路是否打通：

import torch
import torch.nn as nn

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(1, 1)
        # 固定权重：y = 2x + 1
        self.fc.weight.data.fill_(2.0)
        self.fc.bias.data.fill_(1.0)

    def forward(self, x):
        return self.fc(x)

model = LinearModel()
model.eval()
x = torch.randn(1, 1)
torch.onnx.export(model, x, "linear.onnx",
                  input_names=["input"],
                  output_names=["output"],
                  dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})
print("ONNX模型已生成")

3. TensorRT C++推理程序

以下代码完整实现了从ONNX文件构建引擎并执行推理的全过程。将代码保存为 trt_infer.cpp：

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

using namespace nvinfer1;

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

// 读取文件到buffer
std::vector<char> loadFile(const std::string& path) {
    std::ifstream f(path, std::ios::binary);
    return std::vector<char>((std::istreambuf_iterator<char>(f)),
                             std::istreambuf_iterator<char>());
}

// 构建引擎
ICudaEngine* buildEngine(const std::string& onnxPath, int batchSize) {
    auto builder = std::unique_ptr<IBuilder>(createInferBuilder(gLogger));
    auto network = std::unique_ptr<INetworkDefinition>(builder->createNetworkV2(0));
    auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, gLogger));

    auto onnxData = loadFile(onnxPath);
    if (!parser->parse(onnxData.data(), onnxData.size())) {
        std::cerr << "ONNX解析失败" << std::endl;
        return nullptr;
    }

    auto config = std::unique_ptr<IBuilderConfig>(builder->createBuilderConfig());
    config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 1 << 20); // 1MB
    builder->setMaxBatchSize(batchSize);

    return builder->buildSerializedNetwork(*network, *config);
}

int main() {
    const std::string onnxPath = "linear.onnx";
    const int batchSize = 1;

    // 构建引擎
    auto engine = std::unique_ptr<ICudaEngine>(buildEngine(onnxPath, batchSize));
    if (!engine) return -1;

    auto context = std::unique_ptr<IExecutionContext>(engine->createExecutionContext());

    // 获取输入输出binding索引
    int inIdx = engine->getBindingIndex("input");
    int outIdx = engine->getBindingIndex("output");

    // 输入输出数据（batch=1, 特征维度=1）
    float h_in[1] = {3.0f};
    float h_out[1] = {0.0f};

    // 分配GPU显存
    void* d_in;
    void* d_out;
    cudaMalloc(&d_in, sizeof(float));
    cudaMalloc(&d_out, sizeof(float));
    cudaMemcpy(d_in, h_in, sizeof(float), cudaMemcpyHostToDevice);

    // 设置binding地址
    context->setTensorAddress("input", d_in);
    context->setTensorAddress("output", d_out);

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

    // 拷贝结果回主机
    cudaMemcpy(h_out, d_out, sizeof(float), cudaMemcpyDeviceToHost);

    std::cout << "输入: " << h_in[0] << "，推理输出: " << h_out[0] << std::endl;

    // 清理
    cudaFree(d_in);
    cudaFree(d_out);
    return 0;
}

4. 编译与运行

# 编译
g++ -std=c++17 trt_infer.cpp -o trt_infer \
    -I/usr/include/aarch64-linux-gnu \
    -L/usr/lib/aarch64-linux-gnu \
    -lnvinfer -lnvonnxparser -lcudart

# 运行（需与linear.onnx在同一目录）
./trt_infer

验证方法

程序运行后，观察终端输出。由于模型权重固定为 y = 2x + 1，输入3.0时输出应为7.0。判定标准：

程序无报错退出，退出码为0
输出值等于7.0（允许±1e-5浮点误差）
若输出为0或随机值，说明引擎构建或binding设置错误

预期输出：输入: 3，推理输出: 7

常见问题

编译时找不到头文件：确认已安装libnvinfer-dev，并检查-I路径是否正确
ONNX解析失败：确认模型路径正确，且模型为opset 11及以上版本
输出为0：检查是否忘记调用setTensorAddress，或GPU显存未正确分配
首次运行较慢：TensorRT构建引擎需要数秒，属正常现象