首页 / 资讯中心 / 文章详情

TVM Relay 视觉算子 API 全解析:tvm.relay.vision 目标检测算子家族

TVM Relay 视觉算子 API 全解析:tvm.relay.vision 目标检测算子家族 ★ FEATURED ARTICLE
编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载本篇技术指南以 vision.rst 的 API 参考为骨架系统梳理 Apache TVM Relay 中tvm.relay.vision模块提供的视觉算子SSD 的multibox_prior/multibox_transform_loc、目标检测后处理的多种 NMS非极大值抑制、Faster/Mask R-CNN 的roi_align/roi_pool/proposal以及 YOLO/Darknet 的yolo_reorg。读者将掌握每个算子的输入输出格式、全部参数语义与默认值、底层 TOPI 实现与策略注册机制并能在 Relay IR 中正确构造这些算子。tvm.relay.vision 模块概述tvm.relay.vision是 Relay 标准算子库中专门负责计算机视觉目标检测的算子集合通过from tvm import relay后以relay.vision.op_name的形式调用。在 python/tvm/relay/init.py 中vision与image、annotation、contrib、dyn等算子子模块一同被导出。模块的实际定义位于 python/tvm/relay/op/vision/init.py它按检测算法家族聚合了四类实现from .multibox import * # SSD 多框先验生成与位置变换 from .nms import * # 各类非极大值抑制 from .rcnn import * # ROI Align / ROI Pool / Proposal from .yolo import * # YOLO reorg对应的算子注册逻辑分散在_vision.py、_rcnn.py、_yolo.py中。API 文档 vision.rst 通过automoduleautosummary自动从上述 Python 模块的 docstring 提取签名与参数说明因此本文的算子参数细节与仓库源码严格一致。SSD 系列算子先验框生成与位置解码SSDSingle Shot Multibox Detector类检测模型在 Relay 中由两个算子构成先验anchor框生成与位置回归解码。multibox_prior生成先验框relay.vision.multibox_prior(data, sizes(1.0,), ratios(1.0,), steps(-1.0, -1.0), offsets(0.5, 0.5), clipFalse)data输入特征图张量relay.Expr形状为 4-D[batch, channel, height, width]sizesanchor 尺寸元组默认(1.0,)ratiosanchor 宽高比元组默认(1.0,)steps先验框在 y、x 方向的步长-1.0表示自动计算等于特征图尺寸的倒数offsets先验框中心偏移y、x默认(0.5, 0.5)clip是否将越界框裁剪到[0, 1]区间默认False。输出为 3-D 张量形状[1, h_in * w_in * (num_sizes num_ratios - 1), 4]。每个特征图位置生成num_sizes num_ratios - 1个先验框每个框用 4 个归一化坐标描述。multibox_transform_loc解码位置回归relay.vision.multibox_transform_loc(cls_prob, loc_pred, anchor, clipTrue, threshold0.01, variances(0.1, 0.1, 0.2, 0.2), keep_backgroundFalse)cls_prob类别概率张量loc_pred位置回归预测张量anchor由multibox_prior生成的先验框clip是否裁剪越界框默认Truethreshold判定为正样本的置信度阈值默认0.01variances解码框回归输出时使用的方差默认(0.1, 0.1, 0.2, 0.2)keep_background是否保留被判定为背景的框默认False。输出为relay.TupleWrapper包含 2 个元素解码后的框坐标与类别/置信度信息源码见 python/tvm/relay/op/vision/multibox.py。底层实现这两个算子注册在 python/tvm/relay/op/vision/_vision.py均声明为OpPattern.OPAQUE不可融合的复杂算子reg.register_strategy(vision.multibox_prior, strategy.multibox_prior_strategy) reg.register_pattern(vision.multibox_prior, OpPattern.OPAQUE) reg.register_strategy(vision.multibox_transform_loc, strategy.multibox_transform_loc_strategy) reg.register_pattern(vision.multibox_transform_loc, OpPattern.OPAQUE)计算与调度策略在 python/tvm/relay/op/strategy/generic.py 中绑定到 TOPI 实现topi.vision.ssd.multibox_prior与topi.vision.ssd.multibox_transform_loc并配以topi.generic.schedule_multibox_prior/schedule_multibox_transform_loc调度。NMS 系列算子检测后处理核心NMS非极大值抑制用于在重叠的候选框中保留最可信的检测结果。tvm.relay.vision提供了完整的三级 NMS 算子链及两种跨框架兼容实现。get_valid_counts按阈值筛选有效框relay.vision.get_valid_counts(data, score_threshold, id_index0, score_index1)data3-D 张量形状[batch_size, num_anchors, 6]最后一维格式为[class_id, score, box_left, box_top, box_right, box_bottom]无类别时可为 5 维[score, x1, y1, x2, y2]score_threshold有效框的最低分数下限id_index类别所在索引-1表示禁用类别信息score_index分数/置信度所在索引默认1。输出为 3 元组relay.TupleWrappervalid_count1-D每 batch 有效框数量、out_tensor重排后的数据张量有效框被移到顶部、out_indices有效框在原数据中的索引。形状推断由 TVM script 编写见 python/tvm/relay/op/vision/_vision.pyvalid_count形状为(batch,)out_tensor与输入同形out_indices形状为(batch, num_anchors)。non_max_suppressionMXNet 风格 NMSrelay.vision.non_max_suppression( data, valid_count, indices, max_output_size-1, iou_threshold0.5, force_suppressFalse, top_k-1, coord_start2, score_index1, id_index0, return_indicesTrue, invalid_to_bottomFalse)data3-D 张量[batch_size, num_anchors, 6]或[batch_size, num_anchors, 5]可直接使用get_valid_counts的out_tensor输出valid_count有效框数量对应get_valid_counts的valid_count输出indices2-D 张量[batch_size, num_anchors]表示框在原数据中的索引若未经过get_valid_counts其第二维等价于arange(num_anchors)max_output_size每个实例最大输出框数 0时返回全部有效框默认-1iou_thresholdNMS 的 IoU 阈值默认0.5force_suppress是否忽略类别直接抑制所有重叠框默认Falsetop_kNMS 前保留的最大框数-1表示不限制coord_start连续 4 个坐标的起始索引默认2score_index/id_index分数与类别索引return_indices是否返回框索引默认Trueinvalid_to_bottom是否将有效框移到顶部默认False。输出当return_indicesTrue时返回 2-D 张量对索引与有效数量否则返回 3-D 张量[batch_size, num_anchors, 6]或[batch_size, num_anchors, 5]。该算子对应 MXNet 的contrib_vision.nms。all_class_non_max_suppressionONNX/TensorFlow 兼容 NMSrelay.vision.all_class_non_max_suppression( boxes, scores, max_output_boxes_per_class-1, iou_threshold-1.0, score_threshold-1.0, output_formatonnx)boxes3-D 张量(batch_size, num_boxes, 4)scores3-D 张量(batch_size, num_classes, num_boxes)max_output_boxes_per_class每个类别最大输出框数iou_thresholdIoU 测试阈值score_threshold提前过滤低分框的分数阈值output_formatonnx或tensorflow决定输出结构与消费方。该算子逐类别独立执行 NMS对应 ONNX 的NonMaxSuppression与 TensorFlow 的combined_non_max_suppression。输出差异见 python/tvm/relay/op/vision/nms.pyonnx格式2 元组indices形状(batch_size * num_classes * num_boxes, 3)编码 batch、class、box 三个索引按 batch→class→分数降序排列以及标量num_total_detection形状(1,)tensorflow格式3 元组indices形状(batch_size, num_classes * num_boxes, 2)编码 class、box 索引、scores形状(batch_size, num_classes * num_boxes)以及num_total_detection形状(batch_size,)。对应的形状推断函数在 python/tvm/relay/op/vision/_vision.py 中按output_format分支_all_class_nms_shape_func_onnx与_all_class_nms_shape_func_tf。regular_non_max_suppressionTFLite 兼容 NMSrelay.vision.regular_non_max_suppression( boxes, scores, max_detections_per_class, max_detections, num_classes, iou_threshold, score_threshold)boxes3-D 张量(batch_size, num_boxes, 4)坐标编码为(ymin, xmin, ymax, xmax)scores3-D 张量(batch_size, num_boxes, num_classes_with_background)max_detections_per_class每个类别最大输出框数max_detections总的最大输出框数num_classes不含背景的类别数iou_threshold/score_thresholdIoU 与分数阈值。输出为 4 元组detection_boxes(batch_size, max_detections, 4)、detection_classes(batch_size, max_detections)、detection_scores(batch_size, max_detections)、num_detections(batch_size,)对应 TFLite 的 regular NMS。形状推断见_regular_nms_shape_funcpython/tvm/relay/op/vision/_vision.py各输出第二维由attrs.max_detections决定。NMS 家族的策略注册在 python/tvm/relay/op/vision/_vision.py 中四个 NMS 算子全部注册为OpPattern.OPAQUE策略分别绑定vision.get_valid_counts→topi.vision.get_valid_countsvision.non_max_suppression→topi.vision.non_max_suppressionvision.all_class_non_max_suppression→topi.vision.all_class_non_max_suppressionvision.regular_non_max_suppression→topi.vision.regular_non_max_suppression。具体绑定代码见 python/tvm/relay/op/strategy/generic.py三者共用topi.generic.schedule_nms调度而 GPU 侧在 python/tvm/relay/op/strategy/cuda.py 中可提供专用实现。注意get_valid_counts与non_max_suppression都注册了形状函数register_shape_func支持动态形状推断因此在test_any.py等动态 shape 测试中也会被覆盖。Faster / Mask R-CNN 系列算子roi_align双线性插值 ROI 对齐relay.vision.roi_align(data, rois, pooled_size, spatial_scale, sample_ratio-1, layoutNCHW, modeavg)data4-D 特征图[batch, channel, height, width]rois2-D 张量[num_roi, 5]最后一维格式[batch_index, w_start, h_start, w_end, h_end]pooled_size输出尺寸的二元组(ph, pw)spatial_scale特征图尺寸与原始图像尺寸之比等于卷积层总步长的倒数取值范围(0.0, 1.0]sample_ratioROI 对齐的采样率-1默认表示自适应采样layoutNCHW或NHWC默认NCHWmode池化方式支持avg默认与max。输出为 4-D 张量[num_roi, channel, pooled_size, pooled_size]。形状推断函数python/tvm/relay/op/vision/_vision.py按 layout 分支NCHW 输出[num_roi, C, ph, pw]NHWC 输出[num_roi, ph, pw, C]。roi_pool最大池化 ROI 池化relay.vision.roi_pool(data, rois, pooled_size, spatial_scale, layoutNCHW)参数语义与roi_align相同rois格式[batch_index, w_start, h_start, w_end, h_end]spatial_scale为总步长倒数。采用传统最大池化方式计算定义在 python/tvm/relay/op/vision/_rcnn.py 中调用topi.vision.rcnn.roi_pool_nchw。proposalRPN 候选框生成relay.vision.proposal(cls_prob, bbox_pred, im_info, scales, ratios, feature_stride, threshold, rpn_pre_nms_top_n, rpn_post_nms_top_n, rpn_min_size, iou_loss)cls_prob4-D 张量[batch, 2 * num_anchors, height, width]RPN 分类概率bbox_pred4-D 张量[batch, 4 * num_anchors, height, width]框回归预测im_info2-D 张量[batch, 3]格式[im_height, im_width, im_scale]scales/ratiosanchor 窗口的尺度与宽高比元组feature_strideRPN 卷积层中每个单元的感受野大小该层之前所有步长之积thresholdNMS 阈值rpn_pre_nms_top_nNMS 前保留的最高分框数-1表示全部使用rpn_post_nms_top_nNMS 后保留的 proposal 数rpn_min_sizeproposal 的最小高度或宽度iou_loss是否使用 IoU loss。输出为 2-D 张量[batch * rpn_post_nms_top_n, 5]最后一维格式[batch_index, w_start, h_start, w_end, h_end]。完整签名见 python/tvm/relay/op/vision/rcnn.py。ROI 系列的关键注册细节与 SSD/NMS 不同roi_align与roi_pool的 pattern 为OpPattern.OUT_ELEMWISH_FUSABLEpython/tvm/relay/op/vision/_rcnn.py并注册了布局转换register_convert_op_layoutConvertLayoutpass 会将其转换为目标 layout仅支持NCHW/NHWCrois输入必须保持默认 layout。proposal则为OpPattern.OPAQUE策略绑定topi.vision.rcnn.proposalpython/tvm/relay/op/strategy/generic.py。roi_align的策略实现python/tvm/relay/op/strategy/generic.py会根据attrs.layout选择topi.vision.rcnn.roi_align_nchw或 NHWC 版本。YOLO / Darknet 系列算子relay.vision.yolo_reorg(data, stride)yolo_reorg是 Darknet 模型中使用的空间重排层按stride打乱输入张量的值并同时完成形状变换。若输入形状为(n, c, h, w)、步长为s则输出形状为(n, c*s*s, h/s, w/s)。该算子在 python/tvm/relay/op/vision/_yolo.py 中注册为OpPattern.INJECTIVE逐元素注入型可参与算子融合并使用register_injective_schedule标准调度。yolo.py的 docstring 给出了完整的数值示例python/tvm/relay/op/vision/yolo.py输入(1, 4, 2, 2)、stride2时输出(1, 16, 1, 1)通道维按 stride 的棋盘式采样重排。stride1时 reorg 无实际意义。典型调用与测试验证在 Relay IR 中构造视觉算子以 NMS 三件套为例标准用法是先用get_valid_counts过滤再送入non_max_suppressionimport tvm from tvm import relay data relay.var(data, relay.TensorType((1, 1000, 6), float32)) score_threshold relay.const(0.5, float32) valid_count, out_tensor, out_indices relay.vision.get_valid_counts( data, score_threshold, id_index0, score_index1 ) nms_out relay.vision.non_max_suppression( out_tensor, valid_count, out_indices, max_output_size-1, iou_threshold0.5, force_suppressFalse, top_k-1, coord_start2, score_index1, id_index0, return_indicesFalse, )标量参数会在 Python 侧被自动包装为常量例如score_threshold非Expr时经expr.const(score_threshold, float32)转换max_output_size与iou_threshold同理见 python/tvm/relay/op/vision/nms.py这一约定简化了 Relay IR 的手工构造。测试覆盖算子级验证集中在 tests/python/relay/test_op_level5.pytest_multibox_prior第 286 行、test_multibox_transform_loc第 632 行test_get_valid_counts第 392 行、test_non_max_suppression第 435 行、test_all_class_non_max_suppression第 1508 行test_roi_align第 735 行、test_roi_pool第 852 行、test_proposal第 889 行test_yolo_reorg与test_yolo_reorg_infer_shape第 967、982 行。这些测试均以 NumPy 参考实现逐一比对输出是理解各算子精确语义尤其是坐标顺序、return_indices返回值布局、ONNX/TF 输出格式差异的最佳阅读材料。此外roi_align的布局转换逻辑在 tests/python/relay/test_pass_convert_op_layout.py 中有专门用例。小结tvm.relay.vision覆盖了主流目标检测模型所需的全部后端算子SSD 的先验框生成与解码multibox_*、两阶段检测器的 ROI 提取与 RPN proposalroi_*/proposal、跨框架兼容的 NMS 后处理non_max_suppression系列以及 YOLO/Darknet 的yolo_reorg。理解这些算子的参数语义、输出格式与策略注册方式是手工构造检测网络 Relay IR、编写前端转换器或为自定义硬件后端添加支持的基础。需要更深入的实现细节时可沿 python/tvm/relay/op/vision/ → python/tvm/relay/op/strategy/generic.py → TOPIvision计算库 → src/relay/op/vision/nms.cc、rcnn_op.cc、multibox_op.cc的链路逐层下钻。赞分享编译器深度学习模型优化【免费下载链接】tvmOpen deep learning compiler stack for cpu, gpu and specialized accelerators项目地址https://gitcode.com/gh_mirrors/tvm7/tvm点击查看免费下载相关推荐TiXL 图像反馈特效算子家族指南Lib.image.fx.feedback 全部 7 个 Feedback 算子解析TiXL 图像反馈特效算子家族指南Lib.image.fx.feedback 全部 7 个 Feedback 算子解析 本指南以 TiXL开源实时动态图形创音视频图形学桌面应用TVM Relay 算子策略Operator Strategy完全指南从 TOPI compute/schedule 到多目标编译降级TVM Relay 算子策略Operator Strategy完全指南从 TOPI compute/schedule 到多目标编译降级 Relay 算子策编译器深度学习模型优化TVM Relay 图像算子指南tvm.relay.image 模块详解与源码级解析TVM Relay 图像算子指南tvm.relay.image 模块详解与源码级解析 导读 tvm.relay.image 是 Apache TVM Rela编译器深度学习模型优化创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
阅读完成 · 觉得有帮助?
咨询建站