摘要 在深度学习模型部署领域,存在一类特殊的算子,其输出形状即使在 输入形状完全确定 的情况下仍无法在编译期预知。以 NonMaxSuppression (NMS) 为代表的这类算子,给静态图框架的形状推断和内存管理带来了根本性挑战。本文深入剖析此类算子的本质特征,系统梳理现有推理引擎的解决方案,包括 双输出端口 、 输出填充与元数据传递 等机制。进一步,本文提出一种 双形状张量数据结构 的改进思路——在张量中同时维护 最大形状 和 实际形状 ,并探讨基于 统计信息 设置最大输出阈值的优化策略,为框架开发者提供处理运行期动态形状算子的设计参考。
目录 一、静态图框架下的确定性困境 1.1 静态图的基本假设 静态图编译框架(如 TensorRT、OpenVINO、MindSpore 静态图模式)的核心优势在于 编译期优化 。框架在模型加载阶段执行形状推断( infershape ),确定每个张量的形状信息,进而完成内存分配、算子融合、执行流编排等优化。这一过程隐含的基本假设是: 所有张量的形状在编译期可以完全确定 。
1.2 确定性的边界 对于绝大多数算子(卷积、池化、全连接等),输出形状是输入形状和算子参数的 确定性函数 。例如, Conv2d 的输出形状可通过输入形状
与卷积核参数精确计算。这种 静态确定性 构成了传统编译优化的基础。
1.3 运行期动态形状算子的引入 然而,目标检测等高级任务引入了 非极大值抑制(NonMaxSuppression, NMS) 类算子。这类算子的输出数量取决于 数据内容 ——图像中物体的数量、重叠程度、置信度分布,而这些信息在编译期完全未知。这就产生了一类特殊的算子: 输入形状确定,输出形状运行时决定 。
二、运行期动态形状算子的典型特征与分类 2.1 核心特征 运行期动态形状算子具备以下共同特征:
• 输入输出映射的非固定性 :输出张量的第一个维度(通常是检测框数量)无法从输入形状直接推导。 • 数据内容依赖性 :输出数量取决于输入数据的统计特性。 • 阈值参数影响 :算子属性(如 iou_threshold 、 score_threshold )直接调控输出规模。 2.2 典型算子列举 下表整理了常见的运行期动态形状算子:
算子名称 输入描述 输出描述 动态维度 来源框架 NonMaxSuppression boxes[N,4] indices[num_valid,3] num_valid SoftNMS boxes[N,4] dets[num_valid,5] num_valid RoIAlign feat[C,H,W] feat[num_rois,C,outH,outW] num_rois MulticlassNMS boxes[N,4]
dets[num_valid,6] num_valid TopK data[...] values[K] NonZero data[...] indices[num_nonzero, ndim] num_nonzero
2.3 与传统动态形状的区别 需要区分 输入动态形状 与 运行期动态形状 :
维度 输入动态形状 运行期动态形状 变化来源 编译处理 代表算子 Resize NMS 内存策略
2.4 NMS 算子的特点 NMS 作为运行期动态形状的典型代表,具有以下关键特点:
• 输出数量不确定 :即使输入框数量固定,经过置信度筛选和 IoU 抑制后,输出框数量随图像内容变化。
• 阈值敏感性 : score_threshold 和 iou_threshold 直接影响输出规模,调整阈值可显著改变输出数量。 • 排序依赖性 :通常按置信度降序输出,使得截断操作(只取前 K 个)对结果影响相对可控,因为低分框往往贡献较小。 三、端到端部署:包含NMS的完整计算图 将 NMS 算子直接集成到推理模型中,形成 端到端的计算图 ,可以显著简化部署流程。传统做法中,NMS 往往作为模型外部的后处理步骤,由宿主代码实现,这需要额外的开发、调试和优化工作。而将 NMS 纳入计算图后,整个检测流程被封装在一个统一的图结构中,部署时只需加载模型并执行一次推理,即可直接获得最终检测结果。
下图展示了一个典型的目标检测模型(如 YOLO 或 SSD)的完整计算图,其中包含了从图像预处理到 NMS 输出的所有环节:
端到端集成的优势 :
• 部署简化 :推理代码只需调用一次模型执行,无需编写额外的后处理逻辑。 • 统一优化 :框架可对整个图进行联合优化,例如将 NMS 的阈值作为常量折叠、融合相邻操作等。 • 跨平台一致性 :NMS 的具体实现由框架保证,避免不同平台(C++/Python/Java)后处理逻辑的差异。 • 性能提升 :NMS 可在 GPU 上加速执行,减少 CPU-GPU 数据拷贝。 然而,这种集成也带来了本文讨论的核心问题——NMS 的输出形状在编译期未知,必须由框架特殊处理。后续章节将深入探讨各种处理机制。
四、现有推理引擎的解决方案剖析 针对运行期动态形状算子,主流推理引擎发展出多种处理模式。
4.1 双输出端口模式 实现机制 :算子设置两个输出端口——一个输出实际数据,一个输出有效数量。
典型应用 :TensorRT 的 NMS 实现、部分 ONNX Runtime 自定义算子。
# TensorRT NMS输出示例 # output0: [num_detections, 6] - 检测框数据 # output1: [1] - 实际检测框数量(标量) detections, num_detections = engine.execute(bindings) valid_count = num_detections[ 0 ] valid_detections = detections[:valid_count, :] 优点 :元数据与数据分离,清晰明确。 缺点 :需框架支持多输出传递,下游算子需适配。
4.2 输出填充与标记模式 实现机制 :按理论最大容量分配输出张量内存,未使用部分填充特殊标记(如 ),运行时通过扫描确定有效边界。
典型应用 :OpenVINO 的 NonMaxSuppression-3 、昇腾 AI 处理器的 NMS 实现。
# OpenVINO NMS输出示例
# 输出形状: [max_output_boxes_per_class * num_classes, 3] # 有效数据: 连续排列,后续填充[-1, -1, -1] indices = nms_output[:valid_count, :] 从昇腾社区文档中的示例可以看到填充的具体形式:
batch_size = 1, class_num = 2, max_output_boxes_per_class = 5 第一个类筛选出3个框,第二个类筛选出1个框,输出为: [0, 0, 8] # 有效 [0, 0, 5] # 有效 [0, 0, 7] # 有效 [-1, -1, -1] # 填充 [-1, -1, -1] # 填充 [0, 1, 9] # 有效 [-1, -1, -1] # 填充 [-1, -1, -1] # 填充 [-1, -1, -1] # 填充 [-1, -1, -1] # 填充 优点 :输出形状静态化,便于下游算子内存分配。 缺点 :引入无效数据,下游算子需感知有效边界。
4.3 动态内存分配模式 实现机制 :运行时根据实际输出数量动态分配内存,返回精确形状的张量。
典型应用 :ONNX Runtime 官方文档描述的 NMS 实现、PyTorch 动态执行。
# ONNX Runtime NMS输出 # 输出形状: [num_valid_boxes, 3] - 精确大小 indices = session.run([ "selected_indices" ], feed_dict) valid_count = indices.shape[ 0 ] # 直接从形状获取 优点 :使用自然,下游算子直接处理精确形状。 缺点 :内存分配延迟暴露在运行期,可能影响实时性。
4.4 现有方案的局限性 方案 内存效率 使用便捷性 下游算子支持 实时性能 双输出端口 填充标记 动态分配
现有方案在
内存效率 与 实时性能 之间存在权衡,且下游算子往往需要特殊适配才能正确处理填充数据或依赖双输入。
五、双形状张量数据结构的设计构想 针对现有方案的不足,提出一种 双形状张量数据结构 的改进思路,在框架底层统一处理运行期动态形状问题。本设计将所有张量统一为此结构,框架内部仅使用该数据结构,简化算子开发和内存管理。
5.1 核心设计思想 在张量的底层数据结构中同时维护 两个形状信息 :
• 最大形状(Max Shape) :内存分配时采用的形状,基于算子属性(如 max_output_boxes_per_class )计算的理论上限,在 infershape 阶段确定,用于静态内存池分配。 • 实际形状(Actual Shape) :运行时实际包含有效数据的形状,由算子执行后更新,用于确定数据边界和 ND 索引访问。 5.2 简化的数据结构定义 // 统一张量数据结构示意 template < typename T> class Tensor { private : T* data_; // 数据指针 Shape max_shape_; // 最大形状(静态,内存大小,infershape确定) Shape actual_shape_; // 实际形状(动态,运行时更新) public : // 构造函数:基于最大形状分配内存 Tensor ( const Shape& max_shape) : max_shape_ (max_shape), actual_shape_ (max_shape) { data_ = allocator. allocate (max_shape_. num_elements () * sizeof (T)); } // 设置实际形状(由算子运行时调用) void set_actual_shape ( const Shape& actual) { // 实际元素数不能超过最大分配数 assert (actual. num_elements () <= max_shape_. num_elements ()); actual_shape_ = actual; } // 获取实际形状 Shape actual_shape () const { return actual_shape_; } // 获取最大形状 Shape max_shape () const { return max_shape_; } // 获取数据指针(用于底层访问) T* data () { return data_; } const T* data () const { return data_; } // ND索引访问:按照实际形状进行边界检查,返回引用 T& at ( const std::vector< int64_t >& indices) { // 检查维度匹配 assert (indices. size () == actual_shape_. ndim ()); // 检查每个索引不超出实际形状 for ( size_t i = 0 ; i < indices. size (); ++i) {
assert (indices[i] >= 0 && indices[i] < actual_shape_[i]); } // 计算扁平偏移(行优先) size_t offset = 0 ; size_t stride = 1 ; for ( int i = actual_shape_. ndim () - 1 ; i >= 0 ; --i) { offset += indices[i] * stride; stride *= actual_shape_[i]; } return data_[offset]; } // const版本 const T& at ( const std::vector< int64_t >& indices) const { // 类似实现 } }; 设计要点 :
• 最大形状决定内存大小,在 infershape 阶段由框架根据算子参数(如 NMS 的 max_output_boxes_per_class )计算得到,用于内存分配。 • 实际形状决定数据边界,由算子在执行完成后通过 set_actual_shape 设置。 • ND 索引访问通过 at 方法,严格依据 actual_shape_ 进行边界检查,确保不会访问到填充区域。 • 下游算子可优先使用 at 进行安全访问,或通过 data() 和 actual_shape() 自行控制。 5.3 运行期工作流程 5.4 ND索引访问与写入保护机制 当算子需要写入数据时,应使用 at 方法进行 ND 索引写入,以确保只操作实际形状范围内的位置。对于需要连续写入大量数据的场景(如 NMS 填充输出),可先通过 set_actual_shape 告知实际数量,然后使用指针批量操作,但需保证不越界。
// NMS算子实现片段 void NonMaxSuppression::compute (Tensor< float >& boxes, Tensor< float >& scores, Tensor< int64_t >& output_indices) { // 假设 output_indices 的最大形状已设置为 [max_boxes, 3] // 执行NMS算法,得到实际检测框数量 num_det int64_t num_det = perform_nms (boxes, scores, ...); // 填充输出数据 for ( int64_t i = 0 ; i < num_det; ++i) { output_indices. at ({i, 0 }) = batch_id; output_indices. at ({i, 1 }) = class_id; output_indices. at ({i,
2 }) = box_index; } // 设置实际形状 output_indices. set_actual_shape ({num_det, 3 }); } 5.5 内存布局示意图 最大形状: max_shape = [6, 3] (共18个元素) 实际形状: actual_shape = [4, 3] (共12个有效元素) 内存布局 (行优先): +---------+---------+---------+---------+---------+---------+ | row0 | row1 | row2 | row3 | (未使用)| (未使用)| | [0,0,0] | [1,1,1] | [2,2,2] | [3,3,3] | 任意值 | 任意值 | +---------+---------+---------+---------+---------+---------+ ^ ^ | | 有效数据范围 (actual_shape) 超出实际形状的区域 下游算子应仅通过 actual_shape() 获取有效维度,并使用 at 或基于实际形状的循环访问数据。
5.6 优势分析 • 内存静态分配 :保持静态图的内存池优化能力,避免运行时分配开销。最大形状在 infershape 阶段确定,内存池可预先分配。 • 数据边界清晰 :实际形状始终准确,下游算子通过 actual_shape() 即可知有效范围。 • 安全访问 : at 方法提供 ND 索引边界检查,防止意外访问无效区域。 • 统一抽象 :框架内部统一使用 Tensor 类型,算子开发者只需在运行时更新实际形状。 • 灵活性 :支持高性能指针操作与安全访问两种模式。 六、基于统计信息的最大输出阈值设定策略 6.1 理论最大值的局限性 当前 NMS 算子的 max_output_boxes_per_class 属性通常设置为理论上限(如检测数据集的最大可能物体数)。这种做法导致:
• 硬件限制 :某些硬件对输出大小有硬性上限(如昇腾建议不超过 700)。 6.2 统计设定方法论 可以基于 历史数据的统计分析 来设定更合理的最大输出阈值:
1. 数据采集 :在验证集或历史推理数据上收集 NMS 实际输出数量的分布。 6.3 3σ 原则的应用 对于近似正态分布的场景,可采用 3σ 原则 设定上限:
其中 的选择与容忍度相关:
6.4 溢出处理策略与截断影响分析 当实际输出超过统计设定的上限时,可采用以下策略:
• 截断策略 :保留得分最高的前 max_output 个检测框。 对于 NMS 这类算子, 适当的截断对最终算法效果影响通常较小 。原因在于:
• NMS 本身已按置信度排序,保留 top-K 意味着丢弃最低置信度的检测框,而这些框对最终指标(如 mAP)贡献微弱。 • 检测任务中,极端大量输出的情况往往是密集重叠的低分框,截断它们不会引入漏检。 实验数据表明(基于 COCO 验证集),将 max_output_boxes_per_class 从 100 降低到 50,mAP 下降通常 < 0.5%;降低到 30,下降 < 2%。
然而,对于某些对输出完整性敏感的算子(如 NonZero 用于寻找所有非零元素),截断可能导致信息丢失,此时必须设置为理论上限或采用动态分配。因此, 策略选择需结合算子特性 :
七、实践指南与框架适配建议 7.1 算子开发者的最佳实践 开发新的运行期动态形状算子时,建议遵循以下规范:
1. 明确动态维度 :在算子文档中清晰标注哪些输出维度是运行期确定的。 2. 提供容量参数 :为动态输出提供最大容量控制(如 max_output ),该参数将用于 infershape 阶段设置张量的最大形状。 3. 支持实际形状更新 :在算子实现中调用框架的 set_actual_shape 接口。 4. 使用安全访问 :在调试阶段使用 at 方法进行 ND 索引访问;性能关键路径可先用 set_actual_shape 再使用指针操作。 5. 填充值标准化 :如采用填充策略,统一使用 等无效值(但本设计已通过实际形状解决边界问题,可避免填充)。 7.2 框架层面的支持要求 推理引擎若采用双形状设计,需提供以下基础设施:
• 形状传播增强 :编译期形状推断支持基于算子参数的 最大形状 计算,并为每个张量存储最大形状。 • 运行时形状更新 :张量对象的实际形状可在线程安全的方式下更新。 • 内存池适配 :基于最大形状进行内存池管理,确保所有分配都按最大形状预留。 • 算子融合考量 :融合后的算子需正确处理双形状传递,并保证融合后的算子仍能更新实际形状。 • 调试工具 :提供检查张量访问是否越界的工具(例如在 at 中启用断言)。 7.3 下游算子的适配策略 对于可能接收动态形状张量的下游算子,推荐实现方式:
# 下游算子处理动态输入的通用模式 def process_detections ( tensor ): # tensor 是 Tensor 类型 # 获取实际形状 actual_shape = tensor.actual_shape() rows = actual_shape[ 0 ] cols = actual_shape[ 1 ] if len (actual_shape) > 1 else 1 # 安全方式:使用 at 逐元素访问 for i in
range (rows): for j in range (cols): value = tensor.at([i, j]) # 处理 value # 高性能方式:获取指针和实际形状,批量处理 data = tensor.data() for i in range (rows): row_ptr = data + i * cols # 处理这一行 7.4 性能评估指标 引入双形状机制后,应关注以下指标:
• 内存节省率 : 理 论 最 大 内 存 实 际 分 配 内 存 理 论 最 大 内 存 ,其中实际分配内存基于最大形状。 • 有效数据密度 : 实 际 形 状 元 素 数 最 大 形 状 元 素 数 ,衡量内存利用率。 • 形状更新开销 : set_actual_shape 调用的平均耗时。 • 下游算子适配成本 :需修改的算子数量与代码行数。
八、结语与展望 运行期动态形状算子是静态图编译框架面临的 终极不确定性挑战 。本文从 NMS 算子出发,系统分析了此类算子的本质特征,梳理了现有解决方案的优缺点,并提出了 双形状张量数据结构 的设计构想。通过在张量层面同时维护最大形状与实际形状,框架可以在保持静态内存分配优势的同时,提供精确的数据边界和安全高效的 ND 索引访问,实现内存效率、安全性与使用便捷性的统一。
未来研究方向包括:
• 符号形状扩展 :将双形状思想与符号执行结合,实现更智能的形状推断。 • 编译-运行时协同 :基于历史统计信息进行编译期优化(如自动调整最大形状),运行时动态调整。 • 硬件协同设计 :在 NPU 中内置实际形状传递机制,减少主机-设备通信。 • 自动化参数调优 :基于强化学习的 max_output 参数动态调整策略。 • 标准化 :推动 ONNX 等中间表示支持双形状语义,便于不同框架间的模型交换。 随着模型部署场景的日益复杂,运行期动态形状算子的高效支持将成为推理框架的核心竞争力。希望通过本文的探讨,为框架开发者提供有价值的参考,共同推动深度学习部署技术的进步。