从 ONNX 到 TensorRT：模型转换与嵌入式设备部署实战

前言：

在深度学习领域，模型的训练和优化通常在高性能服务器上进行，但随着边缘计算和物联网设备的普及，将这些模型部署到嵌入式设备上以进行实时推理变得越来越重要。对于想要将深度学习模型部署到嵌入式设备上的工程师来说，如何在有限的计算资源下实现高效推理是一个关键挑战。

这时，NVIDIA 的 TensorRT 工具链提供了一套强大的解决方案。通过将 ONNX 模型转换为 TensorRT 引擎，我们可以大幅提升模型的推理性能，并减少对内存和计算资源的需求。TensorRT 支持多种精度（如 FP16 和 INT8），使得在保证准确度的同时能够极大地加速推理速度。

对于使用 NVIDIA 硬件的嵌入式设备，如 Jetson 系列，TensorRT 提供了与硬件紧密集成的优势。它能充分利用 NVIDIA GPU 的并行计算能力以及 CUDA 和 cuDNN 加速库，从而显著提升模型推理速度。相比于其他常规推理框架，TensorRT 可以在不牺牲精度的情况下通过优化模型结构，使得推理速度更快、占用内存更少。

本篇博客将详细介绍如何使用 trtexec 工具将 ONNX 模型转换为 TensorRT 模型，并在嵌入式设备上进行部署，为深度学习工程师提供一条高效部署的路径。

1. 准备 ONNX 模型

首先，你需要一个已经训练好的 ONNX 模型。如果你已经有一个 ONNX 模型，可以直接进行转换。如果没有，可以选择一些常见的深度学习框架（如 PyTorch、TensorFlow）导出 ONNX 模型。例如，使用 PyTorch 的 torch.onnx.export 函数将分类模型转换为 ONNX 格式。

2. 使用 trtexec 工具转换 ONNX 模型

trtexec 是 TensorRT 的一个命令行工具，可以很方便地将 ONNX 模型转换为 TensorRT 引擎文件。你可以通过以下命令将 ONNX 模型转换为 TRT 模型：

trtexec --onnx=your_model.onnx --saveEngine=your_model.trt --fp16

其中，--onnx 参数指定了你的 ONNX 模型文件，--saveEngine 指定了生成的 TensorRT 引擎文件名，--fp16 启用 FP16 精度以加速推理。如果你使用 INT8 量化模型，还需要准备校准数据，并使用 --int8 参数。

3. 部署到板子上

转换完成后，你可以将生成的 .trt 引擎文件部署到你的嵌入式设备上。这里以使用 C++ 编写的推理代码为例，演示如何加载 TensorRT 模型并进行推理。

示例代码讲解

1、日志记录器

TensorRT 需要一个日志记录器来捕获警告和错误信息。你可以通过实现 ILogger 接口来过滤输出。

class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        if (severity != Severity::kINFO) {
            std::cout << "[TensorRT] " << msg << std::endl;
        }
    }
} gLogger;

2、加载模型

使用 nvinfer1::IRuntime 加载从文件中保存的 TensorRT 引擎。

nvinfer1::ICudaEngine* loadEngine(const std::string& engineFile) {
    std::ifstream file(engineFile, std::ios::binary);
    if (!file) {
        std::cerr << "Failed to open engine file: " << engineFile << std::endl;
        return nullptr;
    }

    file.seekg(0, file.end);
    size_t fileSize = file.tellg();
    file.seekg(0, file.beg);

    std::vector<char> engineData(fileSize);
    file.read(engineData.data(), fileSize);

    nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(gLogger);
    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), fileSize);
    runtime->destroy();

    return engine;
}

3、模型初始化和推理

通过 loadEngine 加载模型并创建推理上下文，然后使用 enqueueV2 来进行推理。

// 初始化函数
void initializeModel(const std::string& engineFile) {
    engine = loadEngine(engineFile);
    if (!engine) {
        std::cerr << "Failed to load engine." << std::endl;
        exit(-1);
    }

    context = engine->createExecutionContext();
    if (!context) {
        std::cerr << "Failed to create execution context." << std::endl;
        engine->destroy();
        exit(-1);
    }
}

4、推理流程

在推理过程中，你需要将图像转换为模型输入所需的格式，执行推理并获取结果。

std::vector<float> runInference(cv::Mat& test_img) {
    // 图像预处理和推理代码
}

5、推理结果后处理

获取推理输出后，可以使用 softmax 对分类结果进行处理，这里以分类模型为例，输出最终的分类预测。

std::vector<float> probabilities = runInference(cut_img);
std::vector<std::string> prb = { "0", "1", "2", "3", "4", "5", "6" };
for (size_t i = 0; i < probabilities.size(); ++i) {
    std::cout << prb[i] << ": " << probabilities[i] << std::endl;
}

4. 部署的注意事项

1、导入必要的库

在编写推理代码时，你需要导入 TensorRT 和 CUDA 相关的库。这些库是高效推理的核心，直接决定了模型在 NVIDIA 硬件上的运行效率。通常你需要确保以下库正确安装并配置好：

• TensorRT 库：TensorRT 是模型推理的核心库，负责加载和执行 TensorRT 引擎。
• CUDA：CUDA 是 NVIDIA 提供的 GPU 加速库，用于处理深度学习的并行计算任务。

2、不同设备的兼容性

不同的嵌入式设备可能有不同的硬件规格，例如 NVIDIA Jetson 系列设备（如 Jetson Nano、Jetson Xavier）与服务器端的 NVIDIA GPU 不完全相同。因此，在模型部署前，你可能需要为不同的设备生成专用的 TensorRT 引擎文件。trtexec 在转换过程中，会根据硬件规格（如 CUDA 版本、计算能力）自动调整优化。

具体操作如下：

trtexec --onnx=your_model.onnx --saveEngine=your_model.trt --fp16

3、精度模式选择 

TensorRT 支持多种精度模式（如 FP16、INT8），这些模式可以显著提升推理速度，但并不是所有设备都支持这些精度。FP16 模式在大多数 Jetson 设备上都是支持的，而 INT8 模式则需要额外的校准数据，通常在推理性能和精度之间需要权衡。要使用 INT8 模式，你必须准备一批代表性的数据集用于校准。

trtexec --onnx=your_model.onnx --saveEngine=your_model_int8.trt --int8 --calib=<calibration_data_path>

4、内存管理 

在推理代码中，尤其是在嵌入式设备上，内存资源非常有限，因此你需要特别注意内存的分配和释放。每次推理结束后，确保释放 GPU 上分配的内存，以避免内存泄漏。例如：

cudaFree(d_input);
cudaFree(d_output);

5. 结语

通过使用 trtexec 工具，可以方便地将 ONNX 模型转换为 TensorRT 引擎并部署到嵌入式设备上进行高效推理。通过 C++ 代码，你可以实现模型加载、推理和后处理的全过程。

这篇博客旨在帮助开发者快速了解如何使用 TensorRT 工具链将深度学习模型部署到嵌入式设备，并通过代码示例帮助你轻松上手。