Detectron2 模型部署导出指南TorchScript / ONNX / Caffe2 的完整导出方案【免费下载链接】detectron2Detectron2 is a platform for object detection, segmentation and other visual recognition tasks.项目地址: https://gitcode.com/GitHub_Trending/de/detectron2本指南基于 Detectron2 仓库中的 detectron2.export 模块文档 及其配套的 部署教程、工具脚本 与 C 示例系统讲解如何将 Python 编写的检测模型导出为可部署产物。读完本文你将掌握 tracing、scripting、caffe2_tracing 三种导出方法的适用场景与限制能够使用 export_model.py 一键导出 TorchScript / ONNX / Caffe2 模型并在 Python 或 C 环境下脱离 detectron2 依赖运行推理。一、导出流程的核心概念用 Python 写成的模型要变成可部署的产物必须经过一次导出export过程。围绕这一过程官方文档与源码定义了三个紧密相关的概念导出方法Export Method如何把 Python 模型完整序列化为可部署格式的方式。仓库支持三种tracing、scripting、caffe2_tracing。其中前两者源自 PyTorch 官方的 TorchScript 机制第三种则是 Detectron2 的特有方案——先用 Caffe2 算子替换模型中的部分算子再执行 tracing。格式Format序列化后模型在文件中的描述方式例如 TorchScript、Caffe2 protobuf、ONNX。运行时Runtime加载序列化模型并执行它的引擎例如 PyTorch、Caffe2、TensorFlow、onnxruntime、TensorRT 等。运行时通常与特定格式绑定PyTorch 需要 TorchScript 格式Caffe2 需要 protobuf 格式。在 detectron2/export/init.py 中可以看到detectron2.export包对外暴露的核心 API 包括TracingAdapter、scripting_with_instances、dump_torchscript_IR以及仅在 Caffe2 可用时导入的Caffe2Tracer等接口同时定义了一个稳定的 ONNX opset 版本常量STABLE_ONNX_OPSET_VERSION 11导出 ONNX 时默认使用。三种导出方法的能力对照部署教程中给出了完整的对比矩阵这里完整保留并补充说明项目tracingscriptingcaffe2_tracing支持的格式TorchScriptTorchScriptCaffe2、TorchScript、ONNX支持的运行时PyTorchPyTorchCaffe2、PyTorchC/Python 推理✅✅✅动态分辨率✅✅✅Batch 大小要求固定Constant动态Dynamic不支持 Batch 推理额外运行时依赖torchvisiontorchvisionCaffe2 算子通常已包含在 PyTorch 中Faster/Mask/Keypoint R-CNN✅✅✅RetinaNet✅✅✅PointRend R-CNN✅❌❌Cascade R-CNN✅❌❌需要特别说明的是caffe2_tracing属于即将弃用的方案官方不再计划为其增加对其他格式/运行时的支持但欢迎社区贡献。此外tracing 导出允许动态输入分辨率但输入图片的数量batch size必须固定scripting 则可以支持动态 batch size。二、tracing 与 scripting导出 TorchScript 模型通过 tracing 或 scripting 导出的 TorchScript 模型文件可以在不依赖 detectron2的情况下被 Python 或 C 加载运行。导出的模型通常仍需要 torchvision或其 C 库来提供部分自定义算子。该功能要求PyTorch ≥ 1.8。支持范围Coverage在GeneralizedRCNN与RetinaNet这两个元架构meta architecture下的大多数官方模型tracing 与 scripting 两种模式都支持Cascade R-CNN 与 PointRend 目前仅支持 tracing用户自定义扩展只要本身可被 script 或 trace也同样支持。TracingAdapter让富结构模型可被 tracetorch.jit.trace要求模型的输入输出都是张量元组而 Detectron2 模型的输入输出是dict、Instances、Boxes这类富结构对象。为此 detectron2/export/flatten.py 实现了TracingAdapter它在构造时通过flatten_to_tuple把输入输出拍平flatten成张量元组同时记录一个可序列化的Schemadataclass用于从拍平后的张量重建原始结构Schema体系支持str/bytes、list、tuple、Mappingdict、Instances、Boxes/ROIMasks等类型分别对应IdentitySchema、ListSchema、TupleSchema、DictSchema、InstancesSchema、TensorWrapSchema见 flatten.py通过allow_non_tensor参数可以过滤非张量对象这在只关心单次执行轨迹如统计 FLOPs时很有用但代价是无法再通过 schema 重建输入输出。其核心用法在 flatten.py 的文档字符串中有示例outputs model(inputs) # inputs/outputs 可能是富结构 adapter TracingAdapter(model, inputs) # 现在可以对 adapter 做 trace traced torch.jit.trace(adapter, adapter.flattened_inputs) # trace 模型只能产生拍平后的输出张量元组 flattened_outputs traced(*adapter.flattened_inputs) # adapter 知道如何把拍平结果还原为原始结构 new_outputs adapter.outputs_schema(flattened_outputs)scripting_with_instances解决 Instances 动态属性的难题torch.jit.script要求代码中的属性静态可知而Instances类在 eager 模式下是动态添加属性的难以直接被 scripting 编译。detectron2/export/torchscript.py 中的scripting_with_instances专门解决这一问题创建一个可 script 的new_Instances类行为与Instances类似但所有属性都是静态声明的——属性必须在fields参数中显式列出注册该新类强制 scripting 编译器在编译Instances时使用它。函数执行完会恢复原状因此可以用不同fields分别 script 多个模型。注意该功能只支持 eval推理模式的模型且fields必须包含模型用到的所有属性无论是否是输入/输出目前不支持 detectron2 未定义的数据类型。调用示例from detectron2.export import scripting_with_instances from detectron2.structures import Boxes fields {proposal_boxes: Boxes, objectness_logits: torch.Tensor} scripted_model scripting_with_instances(model, fields)dump_torchscript_IR调试导出模型的利器trace script 导出后dump_torchscript_IR会在指定目录下输出 4 个调试文件model_ts_code.txt逐子模块打印 TorchScript 的 pretty-printed 代码model_ts_IR.txt递归打印所有子模块的 IRmodel_ts_IR_inlined.txt整个图内联后的完整 IRmodel.txtPyTorch 风格的模型结构。这一工具对排查为什么导出后行为不一致非常有用测试用例 tests/test_export_torchscript.py 中也验证了这些文件都能被正常生成且非空。三、caffe2_tracing导出 Caffe2 / ONNX 模型即将弃用Caffe2Tracer定义于 detectron2/export/api.py负责执行 Caffe2-tracing 导出逻辑。它首先把模型的一部分替换为 Caffe2 算子注意部分算子没有 Caffe2 的 GPU 实现并移除后处理逻辑、只保留原始层输出然后把模型导出为 Caffe2、TorchScript 或 ONNX 格式。转换后的模型可以在不依赖 detectron2 / torchvision的情况下在 Python 或 C 中运行于 CPU 或 GPU。它的运行时针对 CPU 与移动端推理做了优化但不适合 GPU 推理。该功能要求ONNX ≥ 1.6。覆盖范围与输入格式支持GeneralizedRCNN、RetinaNet、PanopticFPN三种元架构下的多数官方模型Cascade R-CNN 不支持Batch 推理不支持通过注册机制加入的自定义扩展只要不含控制流或 Caffe2 中不存在的算子例如可变形卷积通常可以开箱即用例如自定义 backbone 和 head。Caffe2Tracer导出的图接收两个输入张量(1, C, H, W)的 floatdata即一张图像通常取值在[0, 255](H, W)通常需要 padding 到 32 的倍数取决于模型架构1×3的 floatim_info每一行是(height, width, 1.0)其中 height、width 是 padding 前的真实图像尺寸。三种导出方法与保存Caffe2Tracer提供三个导出方法from detectron2.export import Caffe2Tracer tracer Caffe2Tracer(cfg, torch_model, inputs) c2_model tracer.export_caffe2() # 返回 Caffe2Model可 save_protobuf() onnx_model tracer.export_onnx() # 返回 onnx.ModelProto ts_model tracer.export_torchscript() # 返回 torch.jit.TracedModule可 .save()构造Caffe2Tracer时传入的inputs是用于 trace 模型的样本输入对于大多数模型随机输入没有检测到目标会导致错误的 trace 结果所以应使用真实图像样本tools/deploy/export_model.py中get_sample_inputs默认从测试集 loader 取第一个 batch。Caffe2Modelapi.py是对 Caffe2 protobuf 格式模型的包装save_protobuf(output_dir)会保存三个文件model.pb图定义可用 Netron 类工具可视化、model_init.pb模型参数、model.pbtxt人类可读的图定义部署时不需要load_protobuf(dir)静态方法用于反向加载需要model.pb和model_init.pbsave_graph(output_file, inputs)可将图保存为 SVG并可记录每个张量的 shape其__call__方法模拟原生 PyTorch 模型的输入输出接口内部自动完成前后处理格式转换底层通过ProtobufDetectionModel实现既可用来与原始 torch 模型做输出对比也可作为实际部署时实现前后处理的参考。注意由于存在 PyTorch/Caffe2 之间的转换开销该方法不适合做性能基准测试。关于 ONNX 导出有一点需要提醒Caffe2Tracer.export_onnx()导出的模型包含仅在 Caffe2 中可用的自定义算子因此不能直接被 onnxruntime 或 TensorRT 等其他运行时执行针对不同运行时做后处理或变换 pass 是可行的但当前仓库不提供支持。四、实战使用 export_model.py 一键导出tools/deploy/export_model.py 是官方提供的完整示例脚本覆盖了三种导出方法、三种格式的组合使用可视为本文所有 API 的落地示范。命令行参数参数可选值默认值说明--formatcaffe2、onnx、torchscripttorchscript输出格式--export-methodcaffe2_tracing、tracing、scriptingtracing导出方法--config-file配置文件路径空模型配置文件--sample-image图片路径None样本输入图片不提供时从测试集取第一个 batch--run-eval布尔开关False导出后用 COCOEvaluator 在测试集上评估转换后模型--output目录路径—转换模型的输出目录opts任意—通过命令行修改配置项追加在末尾典型导出命令将 Mask R-CNN 用 tracing 方法导出为 TorchScriptpython tools/deploy/export_model.py \ --config-file configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \ --output ./output_ts \ --export-method tracing \ --format torchscript用 caffe2_tracing 方法导出为 Caffe2 protobuf并绘制计算图 SVGpython tools/deploy/export_model.py \ --config-file configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml \ --output ./output_c2 \ --export-method caffe2_tracing \ --format caffe2导出 ONNX走 tracing 或 caffe2_tracing 两条路径皆可python tools/deploy/export_model.py \ --config-file configs/COCO-Detection/retinanet_R_50_FPN_3x.yaml \ --output ./output_onnx \ --export-method caffe2_tracing \ --format onnx脚本内部的导出流程export_model.py的main()逻辑如下见 export_model.py解析参数后调用setup_cfg构建配置基于get_cfg()并注册 PointRend 配置项merge_from_file合并配置文件、merge_from_list合并命令行 opts最后freeze()通过build_model(cfg)创建模型用DetectionCheckpointer.resume_or_load(cfg.MODEL.WEIGHTS)加载权重并切到eval()按--export-method分派到export_caffe2_tracing、export_scripting或export_tracing若指定--run-eval用COCOEvaluatorinference_on_dataset在测试集上评估导出后的模型脚本注释提示默认写死为 COCO 评估器其他数据集需自行替换。几个值得注意的实现细节脚本开头设置了torch._C._jit_set_bailout_depth(1)关闭新 shape 的重新特化re-specialization否则--run-eval会很慢export_scripting中定义了一个ScriptableAdapterBase通过不返回 Instances 而返回 dict来绕开 PyTorch 的一个已知问题pytorch/pytorch#46944否则导出的模型无法部署它同时列出了fields字典包含proposal_boxes、objectness_logits、pred_boxes、scores、pred_classes、pred_masks、pred_keypoints、pred_keypoint_heatmaps等属性export_tracing对GeneralizedRCNN使用model.inference(inputs, do_postprocessFalse)作为inference_func保留 ROI mask 输出对RetinaNet则直接调用模型tracing 导出 TorchScript 后脚本构造eval_wrapper手动补上导出模型缺失的最终 resize 步骤即detector_postprocess以便评估时输出格式对齐。支持脚本化评估的 Python 端验证tests/test_export_torchscript.py是上述两种 TorchScript 导出方法的官方测试参照部署教程明确推荐读者先运行这些示例TestScripting对 Mask R-CNN FPN、Mask R-CNN C4、RetinaNet 执行 scripting 导出并用不同尺寸的输入验证动态 batch 与不同形状下 scripted 模型与原始模型输出一致assert_instances_allcloseTestTracing对 Mask R-CNN、Cascade R-CNN、RetinaNet 执行 tracing特意用更小的随机缩放图像做 trace再验证 trace 出的模型在保存、加载、跨设备cpu/cuda迁移后仍能给出与原始模型一致的输出同时检查导出文件中不包含硬编码的设备类型TestTorchscriptUtils验证了flatten_to_tuple/TracingAdapter的 schema 可 JSON 序列化、可反序列化以及dump_torchscript_IR输出文件完整。这些测试精确展示了每种导出方法可复现的最小调用方式是自定义模型导出时最有价值的参考样板。五、C 部署示例与后处理边界tools/deploy/torchscript_mask_rcnn.cpp 提供了一个完整的 C 推理示例演示如何用export_model.py导出的、三种导出方法生成的 TorchScript Mask R-CNN 模型进行推理./torchscript_mask_rcnn model.ts input.jpg EXPORT_METHOD # EXPORT_METHOD 可取 tracing、caffe2_tracing 或 scripting关键实现要点输入构造按导出方法区分tracing 模式需要(C, H, W)的 NCHW 图像张量caffe2_tracing 模式需要(1, C, H, W)的 data 张量外加(1, 3)的 im_info 张量且代码断言图像宽高必须能被 32 整除FPN 模型要求tracing 模式会在图内做 padding而 caffe2_tracing 不会scripting 模式需要Tuple[Dict[str, Tensor]]类型的输入输出解析按导出方法区分tracing 的输出是Instances字段按字母序拍平后的张量元组caffe2_tracing 是遗留顺序的元组scripting 是List[Dict[str, Any]]需按字段名取值运行时设备通过module.buffers()推断整个模型所在设备把输入搬到该设备后执行module.forward并对 CUDA 流做同步后统计推理延迟。配套的 CMakeLists.txt 显示编译该示例需要find_package(Torch)、find_package(OpenCV)以及find_package(TorchVision)tracing/scripting 模式必需C 标准为 C14。后处理导出模型不包含的部分部署教程明确指出转换后的模型不包含把原始层输出加工成格式化预测结果的后处理操作。例如 C 示例最终只产出未经后处理的原始输出28×28 的 mask因为在真实部署中应用往往需要自己实现轻量化的后处理这一步特意留给用户。export_model.py的 tracing 分支也正是因为这个原因才在评估包装器里手动补上detector_postprocessresize 回原图尺寸等操作。对于 Caffe2 格式模型在 Python 中的使用官方提供了Caffe2Model.__call__包装见 api.py它的接口与 PyTorch 版模型完全一致内部自动应用前后处理以匹配格式——既可以作为使用 Caffe2 Python API 的参考也可以作为实际部署中实现前后处理的模板。六、转换到 TensorFlow 的说明除上述方案外tensorpack Faster R-CNN 提供了把少数标准 detectron2 R-CNN 模型转换到 TensorFlow pb 格式的脚本。其原理是翻译配置与权重因此仅支持有限的几个模型适用范围远不如前文三种导出方法。七、选型建议与限制总结需要 TorchScript 动态 batch、追求完整可部署产物优先选择 scripting但要注意 Cascade R-CNN、PointRend 不支持且需要显式声明Instances的全部fields需要 TorchScript 固定 batch、模型包含 Cascade/PointRend选择 tracing输入分辨率可动态变化但输出是拍平后的张量元组需要配合outputs_schema还原结构需要 Caffe2 格式面向 CPU/移动端推理或 ONNX自定义算子不受限的场景使用 caffe2_tracing但它即将弃用、不支持 batch 推理、不含后处理且 ONNX 产物依赖 Caffe2 专属算子所有 TorchScript 导出要求 PyTorch ≥ 1.8caffe2_tracing 要求 ONNX ≥ 1.6三种导出方法共同的前提是模型处于 eval 模式。无论选择哪条路径官方都建议先运行 tests/test_export_torchscript.py 与 tools/deploy/export_model.py 确认示例可复现再针对自己的模型与数据集修改使用脚本化/tracing 的限制需要逐模型适配这也是当前导出流程需要一定用户努力的原因所在。【免费下载链接】detectron2Detectron2 is a platform for object detection, segmentation and other visual recognition tasks.项目地址: https://gitcode.com/GitHub_Trending/de/detectron2创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考