第一章：TensorRT内存模型概述
大家好，我是你们的讲师。今天咱们聊聊TensorRT的内存模型。说实话，很多做推理优化的朋友，一开始都盯着算子融合、精度校准这些“显学”，却忽略了内存管理这个“隐形杀手”。我见过太多项目，模型精度调得不错，一跑起来就显存溢出，或者推理延迟忽高忽低。说白了，问题多半出在内存上。
这一章，我们就从底层开始，把GPU的显存层次、TensorRT的内存池机制，以及内存的生命周期管理，彻底讲透。
1.1 GPU显存层次结构：你离数据有多近？
先问大家一个问题：为什么GPU计算这么快？除了并行核心多，还有一个关键原因——它有精心设计的存储层次。你可以把GPU想象成一个大型工厂：
- 全局内存（Global Memory）：相当于工厂的中央仓库。容量最大（几十GB），但速度最慢。延迟通常在400-800个时钟周期。所有线程都能访问，但每次读写都要走很长的路。
- 共享内存（Shared Memory）：相当于车间里的工作台。容量很小（每个SM通常几十KB到上百KB），但速度极快，延迟只有几个时钟周期。同一个线程块内的线程可以共享。
- 寄存器（Registers）：相当于工人手里的工具。每个线程私有，零延迟。但数量极其有限（每个线程最多255个，实际更少）。
我在项目中遇到过一位同事，把一个大数组频繁从全局内存读取，结果性能惨不忍睹。后来我建议他把热点数据搬到共享内存里，速度直接翻了3倍。你想想看，数据离计算单元越近，效率就越高。这个道理，放到TensorRT里也一样。
核心原则：尽量让数据待在离计算单元更近的地方。全局内存是最后的选择。
1.2 TensorRT内部内存池机制：别让显存“碎片化”
好，理解了硬件层次，我们来看看TensorRT是怎么管理显存的。TensorRT内部维护了一个内存池（Memory Pool）。它不像我们写程序那样，每次需要显存就malloc一次，用完就free。那样会产生大量碎片，而且分配释放的开销很大。
TensorRT的做法是：在构建引擎（build）阶段，就分析出整个推理过程中，所有张量（tensor）的生命周期和大小。然后，它把这些张量“拼”在一起，复用同一块显存。说白了，就是空间换时间，复用减开销。
举个例子：
// TensorRT内部大致是这样做的（伪代码）
// 分析出张量A在时间t1-t2存活，张量B在t3-t4存活
// 如果它们不重叠，就分配同一块显存地址
void* pool_ptr = cudaMalloc(max_size);
// 张量A使用 pool_ptr[0..sizeA]
// 张量B使用 pool_ptr[0..sizeB] // 复用同一块内存
这样做的好处很明显：显存占用大幅降低。我记得有一次优化一个BERT模型，默认情况下显存占用2.1GB，开启内存池复用后，直接降到1.3GB。嗯，这里要注意，这个机制是TensorRT自动完成的，但你可以通过一些配置来影响它。
小技巧：在构建引擎时，可以设置setMemoryPoolLimit来限制内存池的大小。但别设得太小，否则TensorRT会频繁重新分配，反而降低性能。
1.3 内存生命周期管理：从出生到销毁
接下来，我们聊聊内存的生命周期。在TensorRT中，一块显存从分配到释放，大致经历这几个阶段：
- 构建期（Build Time）：TensorRT分析网络结构，计算每个张量的尺寸，规划内存复用方案。此时会分配一些临时内存用于优化。
- 序列化/反序列化（Serialize/Deserialize）：引擎可以保存为文件。注意，序列化时只保存模型结构和权重，不保存运行时内存状态。
- 运行时（Runtime）：创建执行上下文（ExecutionContext），此时才真正分配推理所需的内存池。每次推理，张量在这个池子里“进进出出”。
- 销毁期（Destruction）：销毁引擎和上下文时，TensorRT自动释放所有显存。
我曾经踩过一个坑：在循环里反复创建和销毁ExecutionContext，结果显存泄漏了。后来才发现，每次创建上下文都会分配新的内存池，而销毁时如果没有正确同步，就会残留一些未释放的显存。所以我的建议是：尽量复用ExecutionContext，不要频繁创建销毁。
避坑指南：我曾经在项目中，因为忘记在销毁前调用cudaDeviceSynchronize()，导致显存泄漏。排查了整整两天。记住：销毁任何CUDA资源前，先同步一下。
1.4 实战：如何查看TensorRT的内存使用？
光说不练假把式。我们来看看怎么实际查看TensorRT的内存使用情况。你可以用nvidia-smi，但那个太粗粒度了。更精细的做法是使用TensorRT的profiling工具：
// 开启TensorRT的profiling
config->setProfilingVerbosity(ProfilingVerbosity::kDETAILED);
// 构建引擎后，可以获取每层的显存使用
auto* inspector = engine->createEngineInspector();
std::string layer_info = inspector->getLayerInformation(0, LayerInformationFormat::kJSON);
// 解析JSON，就能看到每层的显存分配情况
我个人习惯在调试阶段，把每层的显存占用打印出来。这样一眼就能看出哪个层是“显存大户”。比如，一个全连接层如果输入输出都很大，它占用的显存可能比卷积层还多。
| 层类型 | 输入尺寸 | 输出尺寸 | 显存占用（MB） |
|---|---|---|---|
| Conv2d | [1, 64, 112, 112] | [1, 128, 56, 56] | 12.5 |
| Linear | [1, 1024] | [1, 512] | 4.0 |
| Attention | [1, 12, 128, 128] | [1, 12, 128, 128] | 24.0 |
看到没？Attention层虽然输入输出尺寸一样，但因为内部有多个中间张量，显存占用反而最高。这就是为什么Transformer模型那么吃显存的原因。
1.5 本章小结
好了，我们来捋一捋。这一章我们讲了三个核心点：
- GPU显存层次：全局内存慢但大，共享内存快但小，寄存器最快但极少。写算子时，尽量用共享内存和寄存器。
- TensorRT内存池：通过复用显存，大幅降低推理时的显存占用。这是TensorRT高效的关键之一。
- 内存生命周期：构建期规划，运行时复用，销毁期清理。记住：复用上下文，销毁前同步。
下一章，我们会深入讲解如何通过配置来精细控制TensorRT的内存使用。比如，怎么设置工作空间大小，怎么用setMemoryPoolLimit来限制显存。到时候我会分享一个实际项目中的调优案例，保证让你收获满满。
好，今天就到这里。有问题欢迎在评论区留言，或者来我的公众号找我。咱们下章见！