社区所有版块导航
Python
python开源   Django   Python   DjangoApp   pycharm  
DATA
docker   Elasticsearch  
aigc
aigc   chatgpt  
WEB开发
linux   MongoDB   Redis   DATABASE   NGINX   其他Web框架   web工具   zookeeper   tornado   NoSql   Bootstrap   js   peewee   Git   bottle   IE   MQ   Jquery  
机器学习
机器学习算法  
Python88.com
反馈   公告   社区推广  
产品
短视频  
印度
印度  
Py学习  »  机器学习算法

深度学习静态图框架中运行期动态形状算子的处理机制与架构设计-以NMS为例

ai算法芯片与系统 • 6 月前 • 156 次点击  

 

摘要

在深度学习模型部署领域,存在一类特殊的算子,其输出形状即使在输入形状完全确定的情况下仍无法在编译期预知。以 NonMaxSuppression (NMS) 为代表的这类算子,给静态图框架的形状推断和内存管理带来了根本性挑战。本文深入剖析此类算子的本质特征,系统梳理现有推理引擎的解决方案,包括双输出端口输出填充与元数据传递等机制。进一步,本文提出一种双形状张量数据结构的改进思路——在张量中同时维护最大形状实际形状,并探讨基于统计信息设置最大输出阈值的优化策略,为框架开发者提供处理运行期动态形状算子的设计参考。

目录

  • • 静态图框架下的确定性困境
  • • 运行期动态形状算子的典型特征与分类
  • • 端到端部署:包含NMS的完整计算图
  • • 现有推理引擎的解决方案剖析
  • • 双形状张量数据结构的设计构想
  • • 基于统计信息的最大输出阈值设定策略
  • • 实践指南与框架适配建议
  • • 结语与展望

一、静态图框架下的确定性困境

1.1 静态图的基本假设

静态图编译框架(如 TensorRT、OpenVINO、MindSpore 静态图模式)的核心优势在于编译期优化。框架在模型加载阶段执行形状推断(infershape),确定每个张量的形状信息,进而完成内存分配、算子融合、执行流编排等优化。这一过程隐含的基本假设是:所有张量的形状在编译期可以完全确定

1.2 确定性的边界

对于绝大多数算子(卷积、池化、全连接等),输出形状是输入形状和算子参数的确定性函数。例如,Conv2d 的输出形状可通过输入形状  与卷积核参数精确计算。这种静态确定性构成了传统编译优化的基础。

1.3 运行期动态形状算子的引入

然而,目标检测等高级任务引入了非极大值抑制(NonMaxSuppression, NMS) 类算子。这类算子的输出数量取决于数据内容——图像中物体的数量、重叠程度、置信度分布,而这些信息在编译期完全未知。这就产生了一类特殊的算子:输入形状确定,输出形状运行时决定

二、运行期动态形状算子的典型特征与分类

2.1 核心特征

运行期动态形状算子具备以下共同特征:

  • • 输入输出映射的非固定性:输出张量的第一个维度(通常是检测框数量)无法从输入形状直接推导。
  • • 数据内容依赖性:输出数量取决于输入数据的统计特性。
  • • 阈值参数影响:算子属性(如 iou_thresholdscore_threshold)直接调控输出规模。

2.2 典型算子列举

下表整理了常见的运行期动态形状算子:

算子名称输入描述输出描述动态维度来源框架
NonMaxSuppressionboxes[N,4]
scores[N]
indices[num_valid,3]num_valid
ONNX, OpenVINO
SoftNMSboxes[N,4]
scores[N]
dets[num_valid,5]
indices[num_valid]
num_valid
MMCV
RoIAlignfeat[C,H,W]
rois[num_rois,5]
feat[num_rois,C,outH,outW]num_rois
MMCV
MulticlassNMSboxes[N,4]
scores[C,N]
dets[num_valid,6]num_valid
TensorRT
TopKdata[...]
k
values[K]
indices[K]
通用
NonZerodata[...]indices[num_nonzero, ndim]num_nonzero
NumPy/PyTorch

2.3 与传统动态形状的区别

需要区分输入动态形状运行期动态形状

维度输入动态形状运行期动态形状
变化来源
输入数据尺寸变化(可变分辨率、变长序列)
算子内部逻辑(筛选、聚类、采样)
编译处理
符号化维度(Any-1
输出维度完全未知,需特殊机制
代表算子Resize
DynamicSlice
NMS
RoIAlignNonZero
内存策略
按最大范围预留
按理论上限预留或动态分配

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{
        // 类似实现

    }
};

设计要点

  • • 所有张量均使用 Tensor 类,统一表示。
  • • 最大形状决定内存大小,在 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. 1. 数据采集:在验证集或历史推理数据上收集 NMS 实际输出数量的分布。
  2. 2. 统计计算:计算均值 、标准差  和分位数。
  3. 3. 阈值设定:根据容忍度选择合适上限。

6.3 3σ 原则的应用

对于近似正态分布的场景,可采用 3σ 原则 设定上限:

其中  的选择与容忍度相关:

 值理论覆盖比例极端值容忍度适用场景
2
95.45%
较低
内存受限设备
3
99.73%
中等
通用场景
4
99.99%
较高
高召回要求场景

6.4 溢出处理策略与截断影响分析

当实际输出超过统计设定的上限时,可采用以下策略:

  • • 截断策略:保留得分最高的前 max_output 个检测框。
  • • 扩展策略:动态扩展内存(如启用后备分配器)。
  • • 降级策略:回退到 CPU 处理溢出部分。

对于 NMS 这类算子,适当的截断对最终算法效果影响通常较小。原因在于:

  • • NMS 本身已按置信度排序,保留 top-K 意味着丢弃最低置信度的检测框,而这些框对最终指标(如 mAP)贡献微弱。
  • • 检测任务中,极端大量输出的情况往往是密集重叠的低分框,截断它们不会引入漏检。

实验数据表明(基于 COCO 验证集),将 max_output_boxes_per_class 从 100 降低到 50,mAP 下降通常 < 0.5%;降低到 30,下降 < 2%。

然而,对于某些对输出完整性敏感的算子(如 NonZero 用于寻找所有非零元素),截断可能导致信息丢失,此时必须设置为理论上限或采用动态分配。因此,策略选择需结合算子特性

算子类型截断影响建议策略
NMS、TopK
较小
可基于统计设置上限
NonZero、Unique
较大
设理论上限或动态分配
RoIAlign
无(RoI数量由输入决定)
输入动态形状处理

七、实践指南与框架适配建议

7.1 算子开发者的最佳实践

开发新的运行期动态形状算子时,建议遵循以下规范:

  1. 1. 明确动态维度:在算子文档中清晰标注哪些输出维度是运行期确定的。
  2. 2. 提供容量参数:为动态输出提供最大容量控制(如 max_output),该参数将用于 infershape 阶段设置张量的最大形状。
  3. 3. 支持实际形状更新:在算子实现中调用框架的 set_actual_shape 接口。
  4. 4. 使用安全访问:在调试阶段使用 at 方法进行 ND 索引访问;性能关键路径可先用 set_actual_shape 再使用指针操作。
  5. 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 等中间表示支持双形状语义,便于不同框架间的模型交换。

随着模型部署场景的日益复杂,运行期动态形状算子的高效支持将成为推理框架的核心竞争力。希望通过本文的探讨,为框架开发者提供有价值的参考,共同推动深度学习部署技术的进步。


Python社区是高质量的Python/Django开发社区
本文地址:http://www.python88.com/topic/193491