4.2 Arena 内存池与复用
4.2 Arena 内存池与复用
本节摘要：ORT 用 Arena/BFC 内存池复用中间张量缓冲，配合 enable_mem_pattern 在静态 shape 下预规划生命周期。ORT-GPU Arena 目标是将跨 session 显存碎片率压到 5% 以下。
模型只要 2GB，ORT 却占 2.8GB
推理时每个算子都要临时中间张量：卷积输出、激活缓存、分块缓冲。如果每次 run 都向设备申请/释放，不仅慢，而且释放后留下的空洞会让显存碎片化——新的大块请求找不到连续空间，只能整体扩容。Arena（内存池）解决的就是这件事：一次向设备申请一大块，内部按需切分与回收，run 之间整块复用。
「模型理论 2GB 显存，ORT 占用 2.8GB」——未启用 mem pattern 或 arena 策略不当导致 碎片与过度预分配。多请求并发时更明显：每个 Session 各持一块 Arena，池子互不共享，碎片率翻倍。2023 年 ORT-GPU 内存池重构强调 跨 EP、跨 session 的零拷贝复用，降低碎片——工业部署需关注 Session 复用而非每请求新建。
Arena 的三个机制
| 机制 | 作用 | 生效前提 |
|---|---|---|
| Arena | 批量分配、按 scope 释放 | 默认开启 |
| mem_pattern | 静态 shape 下预计算张量生命周期 | 静态或可枚举 shape |
| arena_extend_strategy | 控制扩容粒度 | Arena 容量不足时 |
Arena 的核心是"预分配 + 复用"：第一个 run 把张量生命周期跑一遍，之后每次 run 都在已规划的区间里取用，不再触发设备级分配。enable_mem_pattern 开启后，ORT 在静态 shape 下预计算"哪个张量什么时候活、什么时候死"，把活跃区间不重叠的张量叠放到同一块缓冲——这是显存占用从 2.8GB 降到接近 2GB 的关键开关。
配置代码
import onnxruntime as ort so = ort.SessionOptions() so.enable_mem_pattern = True # 静态 shape 下预规划生命周期 so.enable_cpu_mem_arena = True # CPU Arena cuda_opts = { "arena_extend_strategy": "kSameAsRequested", # 按需扩，减少预分配 } sess = ort.InferenceSession( "model.onnx", sess_options=so, providers=[("CUDAExecutionProvider", cuda_opts)], ) # 检查实际显存占用 # 用 torch.cuda.memory_allocated 或 nvidia-smi 对比开启前后
arena_extend_strategy 有两个值：kNextPowerOfTwo（按 2 的幂扩容，预留空间大）与 kSameAsRequested（按请求量扩容，预留小但可能频繁扩）。显存紧张场景用 kSameAsRequested 压预分配，显存充足场景用 kNextPowerOfTwo 减少扩容次数。
动态 shape 的代价
动态 shape 是内存不可预测的主因：mem_pattern 依赖静态生命周期，输入维度一变，模式作废，Arena 退化成"每次都新分配"。代价体现在两个方向——显存占用升高、每次 run 的分配开销回到设备级。工程上不是拒绝动态 shape，而是控制动态维度的档位数：
| 手段 | 效果 |
|---|---|
| batch 档位化（1/4/8/16） | 有限模式，可预分配 |
| 图像尺寸 resize 到固定档 | 减少 shape 组合爆炸 |
| 动态维仅保留业务必需的 | 最小化模式集合 |
内存治理的故障对照
| 现象 | 对策 |
|---|---|
| 显存偏高 | enable_mem_pattern + 单 Session 复用 |
| OOM | kSameAsRequested 减预分配 |
| 多模型 | 每模型独立 Session 或序列化权重 |
| 并发抖动 | 检查多个 Session 的 Arena 是否重叠申请 |
| 首次 run 卡顿 | mem_pattern 首次需要跑一遍规划，属正常 |
💡 固定 shape 服务务必 enable_mem_pattern；动态场景控制 batch 档位数量。这是内存治理的一对总纲：静态就预规划，动态就限档位。
判断直觉与常见误区
⚠️ mem_pattern 对 FP16/INT8 量化模型同样生效，但 shape 变化会静默失效——上线后 batch 从 4 改成 8，显存曲线突变，先查是不是模式失效。
⚠️ Session 复用是内存治理前提。每请求新建 Session 等于每请求重新规划内存，Arena 收益全丢。服务端必须 Session 池化。
本章回顾
- Arena 降分配次数，一次申请整块复用
- mem_pattern 需静态或可枚举 shape
- ORT-GPU Arena 针对碎片，目标 5% 以下
- 动态 shape 是内存不可预测主因，用档位化控制
- Session 复用是内存治理前提
- OOM 先查扩容策略，再查并发 Session 数
显存测量的正确姿势
调内存之前先会量。两个工具各有用途：nvidia-smi 显示进程整体占用，包含驱动与库的固定开销；torch.cuda.memory_allocated 只在同一进程内有效，且不能直接代表 ORT 的分配。ORT 侧更可靠的做法是用 profiling 输出里的内存事件，或对同一 Session 跑多次后看稳定峰值：
import subprocess import onnxruntime as ort sess = ort.InferenceSession("model.onnx", sess_options=so, providers=[("CUDAExecutionProvider", cuda_opts)]) for _ in range(50): sess.run(None, {"input": feed}) # 查看进程显存峰值（Linux） out = subprocess.run( ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"], capture_output=True, text=True) print("gpu mem used MB:", out.stdout.strip())
对比基线要固定输入 shape、固定 batch、固定热身次数。规则是：先跑 50 次热身取稳定峰值，再切换配置重测，差值才是该配置的真实收益。
碎片化的判定与修复
碎片率不直接可见，但有两个间接信号：显存占用随 run 次数缓慢爬升（泄漏或碎片累积）、换一个 batch 后峰值突变。修复路径按序尝试：
| 步骤 | 操作 | 预期 |
|---|---|---|
| 1 | 开启 enable_mem_pattern | 占用贴近理论值 |
| 2 | 固定输入档位 | 模式稳定生效 |
| 3 | kSameAsRequested 扩容 | 峰值下降 |
| 4 | 检查 Session 复用 | 消除重复规划 |
| 5 | 多模型错峰加载 | 避免并发峰值叠加 |
多模型共享显存的规划
一个进程跑多个模型时，内存问题从"单模型峰值"变成"多模型叠加"。两个策略做平衡：串行复用（同一块 Arena 依次服务多个 Session 请求）与并行隔离（每个模型独立 Session，互不干扰）。前者省显存，后者稳延迟，实际部署常按模型优先级混合：
| 模型类型 | 策略 | 原因 |
|---|---|---|
| 主模型 | 独立 Session + 独占 Arena | 稳定延迟 |
| 辅助模型 | 串行复用主模型 Arena | 省显存 |
| 冷门模型 | 按需加载 + 序列化权重 | 不常驻 |
# 串行复用示意：用同一 Session 处理多模型权重切换 # 更常见的做法是 Session 池 + 权重按需加载 import onnxruntime as ort sess = ort.InferenceSession("model.onnx", providers=[("CUDAExecutionProvider", {})]) # 配合 set_providers 或多次 run 复用同一份显存规划
关键约束：共享显存的前提是所有模型共用一个内存规划器。多个独立 Session 各自预分配时，峰值依然会叠加，共享策略就失去意义。规划层面统一入口，执行层面才谈得上复用。
下一节：线程选项与 GPU 的关系。