📚 课程目录第3章模型加载优化

3. 模型加载优化：从ONNX到TensorRT的序列化内存布局

模型加载，说白了就是把训练好的权重从硬盘搬到显存里。但这一步，很多人觉得不就是调个API吗？嗯，没那么简单。我见过不少项目，模型加载就占了上百兆内存，推理时又重复分配，白白浪费资源。

今天咱们就聊聊，从ONNX到TensorRT的序列化过程中，内存到底经历了什么。我个人习惯把这一步叫做“内存布局的二次重构”。

3.1 ONNX的序列化内存结构

ONNX文件本质上是一个Protobuf序列化的二进制文件。它的内存布局，我拆成三块来看：

区域	内容	内存占比
图结构	节点列表、边连接、算子属性	约5-10%
权重数据	卷积核、BN参数、偏置等	约85-90%
元信息	输入输出名称、数据类型、shape	约1-5%

你想想看，一个YOLOv8的ONNX文件大概40-60MB，其中90%都是权重。但这里有个坑——ONNX的权重是平铺存储的，没有任何内存对齐优化。

我曾经踩过的坑：有一次加载一个YOLOv5s的ONNX，发现内存占用比文件大小多了30%。查了半天，原来是Protobuf反序列化时，内部会额外分配临时缓冲区。如果你用Python的onnxruntime加载，它还会再拷贝一份到numpy数组里。一来一回，内存直接翻倍。

3.2 TensorRT的序列化内存布局

TensorRT的序列化引擎（.trt或.engine文件）跟ONNX完全是两码事。它做了三件大事：

权重重排：把权重按GPU最优访存模式重新排列
算子融合：把多个小算子合并成一个大kernel
内存池预分配：提前算好每层需要的显存大小

我画了一张图，帮你理解这个转换过程：

3.3 序列化过程中的内存峰值

这里有个容易被忽视的问题——序列化过程中的内存峰值。我实测过，在把ONNX转成TensorRT引擎时，内存峰值可能是最终引擎文件的3-5倍。

实测数据（YOLOv8s，FP16）：

ONNX文件大小：43MB
TensorRT引擎大小：28MB
序列化过程中峰值内存：约180MB
引擎加载后常驻内存：约35MB

为什么会这样？因为TensorRT在构建引擎时，需要同时持有ONNX的原始数据、中间优化表示、以及最终引擎数据。这三份数据同时存在内存里，峰值自然就上去了。

3.4 内存优化策略

针对序列化过程中的内存问题，我总结了几个实战策略：

策略	做法	效果
流式加载	不一次读入整个ONNX，而是分段解析	减少峰值内存30-50%
权重预量化	在ONNX阶段先做INT8/FP16量化	减少序列化时内存占用
内存池复用	复用临时缓冲区，避免频繁分配释放	减少碎片化，提升速度
延迟加载	只加载需要的子图，不加载整个模型	适合大模型场景

我的个人习惯：在嵌入式设备上（比如Jetson Orin），我会先用trtexec --saveEngine在PC上生成引擎文件，然后直接拷贝到设备上加载。这样设备端就不需要跑序列化过程，内存峰值能降80%。

3.5 代码实战：内存友好的引擎加载

下面是我在实际项目中用的加载方式，注意看内存管理部分：

// 内存友好的TensorRT引擎加载
class EngineLoader {
private:
    std::vector<char> engineData;
    nvinfer1::IRuntime* runtime;
    nvinfer1::ICudaEngine* engine;

public:
    bool loadEngine(const std::string& path) {
        // 1. 获取文件大小，避免一次性读入大文件
        std::ifstream file(path, std::ios::binary | std::ios::ate);
        if (!file.is_open()) return false;

        size_t fileSize = file.tellg();
        file.seekg(0, std::ios::beg);

        // 2. 预分配固定大小缓冲区，避免多次扩容
        engineData.resize(fileSize);
        
        // 3. 分块读取，每块4KB
        const size_t chunkSize = 4096;
        size_t offset = 0;
        while (offset < fileSize) {
            size_t toRead = std::min(chunkSize, fileSize - offset);
            file.read(engineData.data() + offset, toRead);
            offset += toRead;
        }
        file.close();

        // 4. 创建runtime并反序列化
        runtime = nvinfer1::createInferRuntime(gLogger);
        engine = runtime->deserializeCudaEngine(
            engineData.data(), engineData.size());

        // 5. 反序列化完成后立即释放原始数据
        engineData.clear();
        engineData.shrink_to_fit();

        return engine != nullptr;
    }
};

注意：上面的shrink_to_fit()很关键。我见过有人忘了这步，导致std::vector保留了已释放的内存空间，白白多占了几十MB。C++的vector不会自动缩容，你得手动告诉它。

3.6 序列化格式对比

最后，我把常见的序列化格式做个对比，方便你选型：

格式	加载速度	内存占用	适用场景
ONNX	慢（需解析Protobuf）	高（平铺存储）	模型交换、跨框架
TensorRT（FP32）	快（二进制加载）	中（有优化）	GPU推理、高性能
TensorRT（FP16）	快	低（权重减半）	边缘设备、Jetson
TensorRT（INT8）	快	最低（1/4 FP32）	极致性能、精度可接受

嗯，到这里你应该明白了——模型加载不只是调个API那么简单。从ONNX到TensorRT，内存布局经历了从“描述型”到“执行型”的转变。理解这个过程，你才能写出内存友好的部署代码。