【Bug已解决】[Performance] Large CPU memory regression starting in 1.24 (Java, macOS arm64) for fully-convolutional models with dynamic input shapes 解决方案一、现象长什么样升级到 ONNX Runtime 1.24 后在 macOS arm64 的 Java 应用里跑全卷积模型fully-convolutional如分割/超分网络且输入形状每次都变化dynamic input shapes发现常驻内存比以前版本暴涨数倍且推理越跑越慢不是泄漏是每次推理都临时申请一大块 workspace 又不复用。现象# 现象 A内存随不同输入形状“阶梯式”增长后稳定在高点 # 每换一种输入分辨率就新申请一份卷积 workspace旧的没回收/没复用 # 现象 B同一形状重复跑不增长但换形状就涨 # 说明 workspace 是按“形状”缓存的但 1.24 的缓存键或复用逻辑坏了 # 每次新形状都新建且不淘汰旧形状 # 现象 C只在 Java macOS arm64 全卷积 动态形状 组合触发 # 固定形状正常非卷积transformer正常其他平台内存增长不明显最坑的是现象 A/B不是传统泄漏同形状不涨而是“按形状缓存的 workspace 无限堆积”——跑过 50 种分辨率就留 50 份 workspace移动端/嵌入式内存吃不消。二、背景ONNX Runtime CPU EP 跑卷积时会为某些算子如 Winograd、大 kernel im2col分配一块临时workspace中间缓冲区。为了不每次推理都重新分配ORT 维护一个“按输入形状缓存 workspace”的池子见到新形状就分配并缓存见到旧形状就复用。1.24 引入的某个改动与 Java 绑定或 arm64 的内存分配器调整相关破坏了 workspace 缓存的复用/淘汰逻辑要么缓存键不再命中每次都当新形状要么旧形状的 workspace 没被放回池子复用而是泄漏式堆积。结果是“每个不同形状都独占一份 workspace且互不复用”内存随见过的形状数线性增长。这是性能/内存审查里典型的坑资源池的缓存键或回收逻辑被破坏导致本该复用的缓冲被重复分配。三、根因workspace 缓存键失效1.24 改了形状到缓存键的序列化可能受 Java 侧 shape 数组表示影响导致同一形状算出的键不同缓存永不命中 → 每次新分配。旧 workspace 未回收/未复用分配路径没把用过的 workspace 归还池子或淘汰策略缺失旧形状 workspace 堆积。缺少动态形状的内存回归测试CI 只跑固定形状动态形状的内存曲线从不被断言回归长期存在。本质是workspace 池的缓存键/回收逻辑在 1.24 被破坏导致不同形状各自独占 workspace 不复用且缺动态形状内存测试。四、最小可运行复现下面用 Python 模拟“缓存键失效导致不同形状各自独占 workspace 堆积”class WorkspacePoolBuggy: def __init__(self): self._pool {} # key - workspace self._alloc_count 0 def key_of(self, shape): # BUG: 把 shape 当成 list 和 tuple 混用键不一致 if isinstance(shape, tuple): return str(shape) return str(list(shape)) # tuple 和 list 同形状得到不同 key def get_workspace(self, shape): key self.key_of(shape) if key not in self._pool: self._pool[key] object() # 分配一份 self._alloc_count 1 return self._pool[key] pool WorkspacePoolBuggy() # 同一形状因传入 tuple/list 不同被当成两个键 pool.get_workspace((1, 3, 64, 64)) pool.get_workspace([1, 3, 64, 64]) print(alloc_count (buggy):, pool._alloc_count) # 2本应是 1 class WorkspacePoolFixed: def __init__(self): self._pool {} self._alloc_count 0 def key_of(self, shape): return str(tuple(shape)) # 统一规范化为 tuple def get_workspace(self, shape): key self.key_of(shape) if key not in self._pool: self._pool[key] object() self._alloc_count 1 return self._pool[key] pf WorkspacePoolFixed() pf.get_workspace((1, 3, 64, 64)) pf.get_workspace([1, 3, 64, 64]) print(alloc_count (fixed):, pf._alloc_count) # 1正确复用buggy同形状分配 2 次fixed复用为 1 次。五、解决方案第一层最小直接修复最小修复把形状规范化为稳定的缓存键并保证用过的 workspace 归还复用class WorkspacePoolFixed: def __init__(self, max_cached16): self._pool {} self._max_cached max_cached def _key(self, shape): return str(tuple(int(s) for s in shape)) # 统一规范化 def get_workspace(self, shape): key self._key(shape) if key not in self._pool: self._pool[key] self._allocate(shape) # 简单 LRU 淘汰防止无限堆积 if len(self._pool) self._max_cached: self._pool.pop(next(iter(self._pool))) return self._pool[key] def _allocate(self, shape): return bytearray(1024) # 示意真实是卷积 workspace这一层改动最小规范键 上限淘汰复用恢复、堆积消失。但依赖“每处池都写对”下看第二层。六、解决方案第二层结构性改进把“workspace 池的键规范化 复用 上限淘汰”固化成单一事实来源。下面这个 dataclass 集中管理缓存契约from dataclasses import dataclass, field from typing import Dict, Tuple, Any dataclass class OrtDynamicShapeMemPolicy: 单一事实来源动态形状 workspace 池的缓存与复用契约。 max_cached: int 16 _pool: Dict[Tuple[int, ...], Any] field(default_factorydict) alloc_count: int 0 staticmethod def normalize(shape) - Tuple[int, ...]: return tuple(int(s) for s in shape) # 稳定键tuple of int def get(self, shape, allocator): key self.normalize(shape) if key not in self._pool: self._pool[key] allocator(shape) self.alloc_count 1 # 上限淘汰防止按形状无界堆积 if len(self._pool) self.max_cached: self._pool.pop(next(iter(self._pool))) return self._pool[key] def assert_reuse(self, shape, allocator) - None: 同一形状第二次获取不应新增分配。 before self.alloc_count self.get(shape, allocator) self.get(shape, allocator) if self.alloc_count ! before 1: raise AssertionError(workspace not reused for same shape)这一层的关键收益稳定键normalize把 tuple/list/int 都规范成tuple[int]缓存必命中上限淘汰超过max_cached自动淘汰最旧杜绝按形状无界堆积复用断言assert_reuse验证同形状不重复分配单一事实来源所有 workspace 池约定收口在OrtDynamicShapeMemPolicy。七、解决方案第三层断言 / CI 守护把第二层钉成 pytest挂进 CI覆盖动态形状内存回归import pytest from your_package.ort_dyn_shape_mem import OrtDynamicShapeMemPolicy def _alloc(shape): return bytearray(64) def test_same_shape_reused(): # 断言 1同形状tuple/list 混用只分配一次 p OrtDynamicShapeMemPolicy() p.assert_reuse((1, 3, 64, 64), _alloc) p.get([1, 3, 64, 64], _alloc) # 再次确认 list 也复用 assert p.alloc_count 1 def test_distinct_shapes_cached_but_bounded(): # 断言 2不同形状各自缓存但不超上限 p OrtDynamicShapeMemPolicy(max_cached4) for i in range(20): p.get((1, 3, 32 i, 32 i), _alloc) assert len(p._pool) 4 def test_no_unbounded_growth(): # 断言 3跑过 N 种形状池大小仍受 max_cached 约束 p OrtDynamicShapeMemPolicy(max_cached8) for i in range(100): p.get((1, 1, i 1, i 1), _alloc) assert len(p._pool) 8 def test_key_stable_across_types(): # 断言 4tuple 和 list 同形状键一致 p OrtDynamicShapeMemPolicy() assert p.normalize((1, 2)) p.normalize([1, 2])四条断言从“同形状复用”“不同形状有界”“无界增长防护”“键跨类型稳定”四面把内存回归钉死在 CI。八、排查清单ONNX Runtime 升级后动态形状卷积内存暴涨时同形状重复跑是否增长不增长但换形状就涨 → workspace 按形状缓存且复用坏现象 B。workspace 池的缓存键是否跨 tuple/list/int 表示一致不一致就永不命中每次新分配。是否有上限淘汰没有就按见过形状数无界堆积现象 A。用第二层OrtDynamicShapeMemPolicy键规范化 复用 上限淘汰。加第三层 pytest断言“同形状复用、不同形状有界、无界增长防护、键跨类型稳定”。动态形状的内存曲线必须进 CI 回归固定形状测试发现不了这类问题。九、小结ORT 1.24 的动态形状卷积内存回归本质是workspace 池的缓存键在 Java/arm64 下不再稳定tuple/list/int 表示混用、且缺上限淘汰导致不同形状各自独占一份 workspace 不复用、无界堆积内存随见过的分辨率数线性暴涨。修复分三层——第一层把形状规范化为稳定键并加 LRU 上限淘汰第二层用OrtDynamicShapeMemPolicy这个 dataclass 把“键规范化 复用 上限淘汰”收口成单一事实来源并加复用断言第三层用四条 pytest 把“同形状复用、不同形状有界、无界增长防护、键跨类型稳定”钉死在 CI。核心心法按形状缓存的资源池键必须规范化到唯一表示且必须有上限淘汰否则动态形状输入会按形状数无界堆积内存。