TensorFlow LSTM Object Detection 模型导出 TFLite 完整指南：从 checkpoint 到 FlatBuffer 的两步转换

TensorFlow LSTM Object Detection 模型导出 TFLite 完整指南从 checkpoint 到 FlatBuffer 的两步转换【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models导读LSTDLSTM SSD视频目标检测模型在训练完成后若要部署到移动端 / 嵌入式 TFLite 运行时需要把 TensorFlow 训练出的 checkpoint 依次转换成「TFLite 兼容的冻结图」与「TFLite FlatBuffer 模型」两个中间产物。本文基于 exporting_models.md 的官方流程结合 models 仓库 中export_tflite_lstd_graph.py、export_tflite_lstd_model.py及其底层库实现完整讲解每一步的命令、参数含义、图结构与节点规格并给出使用test_tflite_model.py校验导出的模型是否可正常推理的收尾方法。读完本文你将能够独立把一个训练好的 LSTD checkpoint 导出为可在应用侧直接加载的*.tflite文件并理解冻结图输入 / 输出张量为何被设计成视频序列形式。为什么 LSTD 的 TFLite 导出需要两步LSTD 模型包含 LSTM 时序单元其输入不是单张静态图片而是一段长度为unroll_length的视频帧序列维度为[unroll_length, height, width, 3]。这就决定了它与普通 SSD 模型的导出路径不同不能直接走标准的 saved_model / frozen graph 单步转换。从 exporting_models.md 可以提炼出官方推荐的固定两步流水线从 checkpoint 导出 TFLite 兼容的冻结图frozen graph——把训练权重固化成一张输入 / 输出已按 TFLite 约定定义的tflite_graph.pb从冻结图导出 TFLite FlatBuffer 模型——把上一步的tflite_graph.pb交给 TFLite 转换器产出最终供应用加载的*.tflite文件。在 research/lstm_object_detection 目录下这两个步骤分别对应可执行脚本 export_tflite_lstd_graph.py 与 export_tflite_lstd_model.py。下面逐一展开。导出前的工程与配置前提运行目录与依赖导出命令设计为从tensorflow/models/research目录即本仓库的 research 目录发起执行命令行中以lstm_object_detection/xxx.py的相对模块路径调用脚本底层实现 export_tflite_lstd_graph_lib.py 除导入lstm_object_detection自身外还依赖同仓库 research/object_detection 下的exporter、builders.graph_rewriter_builder、builders.post_processing_builder、core.box_list等模块因此需要保证该目录位于PYTHONPATH中脚本基于tensorflow.compat.v1编写见两处脚本的 import 与tf.app.run请使用兼容 TensorFlow 1.x API 的 TensorFlow 版本环境运行。pipeline config 中的硬性约束从源码实现export_tflite_lstd_graph_lib.py可以确认参与导出的 pipeline config 必须满足以下条件否则会直接抛错约束说明源码行为检测模型类型只能是 SSD仅model { ssd { ... } }支持 TFLite 导出非 SSD 时抛出ValueError(Only ssd models are supported in tflite. ...)图像缩放器必须是fixed_shape_resizerTFLite 侧需要完全确定的高度 / 宽度其它 resizer 类型抛出ValueErrorlstm_model配置块必须存在需要其中的eval_unroll_length决定输入序列长度导出时读取lstm_config.eval_unroll_length输入必须是固定尺寸的 RGB或灰度视频序列输入占位符为[eval_unroll_length, height, width, 3]若开启convert_to_grayscale则通道数为 1依 config 动态构造 placeholder仓库给出的两份参考配置可以印证这些字段的实际写法lstm_ssd_mobilenet_v1_imagenet.configtrain_unroll_length: 4、eval_unroll_length: 4image_resizer { fixed_shape_resizer { height: 256 width: 256 } }检测类别数num_classes: 30使用 Faster R-CNN box codery_scale/x_scale: 10.0、height_scale/width_scale: 5.0后处理采用batch_non_max_suppression与score_converter: SIGMOIDlstm_ssd_interleaved_mobilenet_v2_imagenet.config基于 Interleaved MobileNet V2 的 LSTD 变体image_resizer为 320×320同样配置了eval_unroll_length: 4。注意 TFLite 导出读取的是eval 路径下的固定分辨率与eval_unroll_length这些字段最终决定了冻结图的输入张量形状训练与导出配置需保持一致。第一步从 checkpoint 导出 TFLite 冻结图官方命令在 exporting_models.md 中官方给出从tensorflow/models/research目录执行如下命令# 位于 research 目录下执行 PIPELINE_CONFIG_PATH{path to pipeline config} TRAINED_CKPT_PREFIX/{path to model.ckpt} EXPORT_DIR{path to folder that will be used for export} python lstm_object_detection/export_tflite_lstd_graph.py \ --pipeline_config_path ${PIPELINE_CONFIG_PATH} \ --trained_checkpoint_prefix ${TRAINED_CKPT_PREFIX} \ --output_directory ${EXPORT_DIR} \ --add_preprocessing_op执行成功后${EXPORT_DIR}目录下将出现两个文件tflite_graph.pb二进制格式的 TFLite 兼容冻结图tflite_graph.pbtxt同一张图的文本格式版本便于人工检查节点结构。脚本完整参数表以源码为准对照 export_tflite_lstd_graph.py 中的 flag 定义可以拿到比文档更完整的参数语义与默认值。注意文档命令行中写作--add_preprocessing_op而当前仓库源码中对应的 flag 实际名为add_postprocessing_op且默认值为True历史版本命名可能有差异请以你所 checkout 的源码为准。Flag类型默认值必填含义--output_directorystringNone是导出产物输出目录--pipeline_config_pathstringNone是pipeline_pb2.TrainEvalPipelineConfig格式的配置文件路径--trained_checkpoint_prefixstringNone是训练 checkpoint 前缀如.../model.ckpt--max_detectionsinteger10否最多输出的检测框数量--max_classes_per_detectioninteger1否每个检测框最多输出的类别数--detections_per_classinteger100否Regular NMS 中每个类别使用的 anchor 数量--add_postprocessing_opboolTrue否是否在图中追加 TFLite 自定义后处理算子--use_regular_nmsboolFalse否后处理算子使用 Regular NMS 还是 Fast NMS--config_overridestring否覆盖 pipeline config 指定字段的文本 proto导出图的输入与输出规格根据脚本与库实现导出的tflite_graph.pb中节点约定如下也是后续第二步转换时引用张量名的依据输入节点input_video_tensorfloat32张量形状[unroll_length, height, width, 3]内容是归一化后的输入视频帧序列其中height/width必须与 pipeline config 的fixed_shape_resizer中配置一致unroll_length取自lstm_model.eval_unroll_length。输出节点未附加后处理时即add_postprocessing_opFalseraw_outputs/box_encodingsfloat32形状[1, num_anchors, 4]编码后的框回归预测raw_outputs/class_predictionsfloat32形状[1, num_anchors, num_classes]经过 score conversionSIGMOID 等后的每 anchor 类别得分anchorsfloat32常量张量形状[num_anchors, 4]中心点 / 尺寸编码后的 anchor 集合。输出节点附加后处理时即add_postprocessing_opTrue图中会追加名为TFLite_Detection_PostProcess的自定义算子节点其四个输出为detection_boxesfloat32形状[1, num_boxes, 4]框坐标detection_classesfloat32形状[1, num_boxes]类别索引detection_scoresfloat32形状[1, num_boxes]类别得分num_boxes大小为 1 的float32张量表示检测到的框数量。底层做了什么冻结 追加 TFLite 自定义算子深入 export_tflite_lstd_graph_lib.py 可以看到导出函数的完整调用链读取并校验配置config_util.get_configs_from_pipeline_file解析 pipeline config取出model、lstm_model、eval_config并校验 SSD 与fixed_shape_resizer收集解码与 NMS 所需参数从batch_non_max_suppression读取score_threshold/iou_threshold从faster_rcnn_box_coder读取y_scale/x_scale/h_scale/w_scale从model.ssd.num_classes读取类别数export_tflite_lstd_graph_lib.py构造视频输入与推理图创建形状为[eval_unroll_length, height, width, 3]的input_video_tensor占位符调用model_builder.build(...)构建检测模型并执行 preprocess / predictexport_tflite_lstd_graph_lib.pyNMS 被有意推迟到 TFLite 自定义算子中完成固化原始输出在raw_outputs命名空间下输出box_encodings、class_predictions与常量anchorsanchor 通过get_const_center_size_encoded_anchors转成常量节点若配置中存在graph_rewriter还会在此处执行量化图改写若eval_config.use_moving_averages为真则先将滑动平均变量替换进图再冻结export_tflite_lstd_graph_lib.py冻结图使用exporter.freeze_graph_with_def_protos将 checkpoint 权重固化输出节点为raw_outputs/box_encodings、raw_outputs/class_predictions、anchors按需追加后处理算子若add_postprocessing_opTrue调用append_postprocessing_op把TFLite_Detection_PostProcess节点挂到raw_outputs/*与anchors之上并通过TransformGraph的strip_unused_nodes变换剪掉无用节点export_tflite_lstd_graph_lib.py。该节点的属性max_detections、nms_score_threshold、nms_iou_threshold、y_scale、x_scale等全部来自第 2 步收集的配置落盘分别写出tflite_graph.pb与tflite_graph.pbtxt。其中第 6 步正是 TFLite 目标检测模型常见的做法把 NMS 等后处理做成TFLite custom op在端上执行从而避免 TFLite 转换器不支持这些算子的兼容性问题。用 config_override 微调推理行为而不改训练配置脚本还提供了--config_override参数允许在不改动训练 / 评估配置的前提下对导出的推理图做小规模参数覆盖。它同样接受pipeline_pb2.TrainEvalPipelineConfig文本 proto源码 docstring 中给出的例子是把 NMS 的iou_threshold改为 0.5、score_threshold改为 0.0python lstm_object_detection/export_tflite_lstd_graph.py \ --pipeline_config_path path/to/lstm_pipeline.config \ --trained_checkpoint_prefix path/to/model.ckpt \ --output_directory path/to/exported_model_directory \ --config_override \ model{ \ ssd{ \ post_processing { \ batch_non_max_suppression { \ score_threshold: 0.0 \ iou_threshold: 0.5 \ } \ } \ } \ } \ 可以看到该机制在导出「低阈值召回更多框」或「收紧重叠抑制」等推理变体时非常实用。第二步从冻结图导出 TFLite FlatBuffer 模型官方命令拿到第一步的tflite_graph.pb后在tensorflow/models/research目录执行# 位于 research 目录下执行 FROZEN_GRAPH_PATH{path to exported tflite_graph.pb} EXPORT_PATH{path to filename that will be used for export} PIPELINE_CONFIG_PATH{path to pipeline config} python lstm_object_detection/export_tflite_lstd_model.py \ --export_path ${EXPORT_PATH} \ --frozen_graph_path ${FROZEN_GRAPH_PATH} \ --pipeline_config_path ${PIPELINE_CONFIG_PATH}三个参数均为必填见脚本中的mark_flag_as_required。执行成功后${EXPORT_PATH}即为可直接交付给移动应用加载的 FlatBuffer 模型文件。脚本内部如何调用转换器对照 export_tflite_lstd_model.py 的main实现转换逻辑如下configs config_util.get_configs_from_pipeline_file(FLAGS.pipeline_config_path) lstm_config configs[lstm_model] input_arrays [input_video_tensor] output_arrays [ TFLite_Detection_PostProcess, TFLite_Detection_PostProcess:1, TFLite_Detection_PostProcess:2, TFLite_Detection_PostProcess:3, ] input_shapes { input_video_tensor: [lstm_config.eval_unroll_length, 320, 320, 3], } converter tf.lite.TFLiteConverter.from_frozen_graph( FLAGS.frozen_graph_path, input_arrays, output_arrays, input_shapesinput_shapes) converter.allow_custom_ops True tflite_model converter.convert()几个值得注意的细节输入张量为第一步冻结图唯一的输入input_video_tensor输出张量为后处理自定义算子的 4 个输出端口TFLite_Detection_PostProcess及其:1、:2、:3分别对应detection_boxes / detection_classes / detection_scores / num_boxes——这要求第一步导出时必须开启后处理算子否则冻结图不存在该节点输入形状中的序列长度取自lstm_model.eval_unroll_length而当前仓库源码中高 / 宽硬编码为320, 320因此该流程面向 320×320 固定输入分辨率设计若训练配置使用其它分辨率需要相应调整此处源码这也侧面印证了导出的图对固定形状的强约束converter.allow_custom_ops True是关键开关由于图中包含TFLite_Detection_PostProcess这一自定义算子必须允许 custom ops否则转换器会因遇到未知算子而报错转换完成后直接以二进制方式写入${EXPORT_PATH}。收尾验证用随机输入跑通导出的 tflite仓库在 test_tflite_model.py 中提供了面向 TFLite 模型的冒烟测试工具可用于快速确认 FlatBuffer 可被解释器加载并完成一次前向推理python lstm_object_detection/test_tflite_model.py \ --model_path ${EXPORT_PATH}该脚本会依次执行通过tf.lite.Interpreter(model_path...)加载模型并allocate_tensors()打印输入 / 输出张量的input_details与output_details其中会显示输入形状可据此核对是否为[unroll_length, 320, 320, 3]用np.random.random_sample生成随机输入喂入网络并invoke()打印第一个输出张量的内容。脚本 docstring 也提示随机输入仅用于验证实际使用时应把input_data替换为按同形状组织好的真实归一化视频帧。若该步能正常打印输出说明两步导出产出的 tflite 模型在运行时层面是完整的。小结与常见问题排查回顾整条链路从 checkpoint 到可交付的 TFLite FlatBuffer 共两步、对应两个脚本阶段脚本输入输出导出冻结图export_tflite_lstd_graph.pypipeline config checkpointtflite_graph.pb/tflite_graph.pbtxt导出 TFLiteexport_tflite_lstd_model.pytflite_graph.pb pipeline config${EXPORT_PATH}FlatBuffer冒烟验证test_tflite_model.pytflite 模型路径前向推理输出排障时可以优先对照以下几点均可从前述源码与配置文件中找到依据pipeline config 必须是 SSD 模型且使用fixed_shape_resizer同时存在[lstm_object_detection.protos.lstm_model]配置块参考两个示例 config输入分辨率与eval_unroll_length在训练、pipeline config、导出脚本三处必须自洽若应用侧只消费检测框结果detection_boxes等 4 项输出第一步必须保留add_postprocessing_opTrue默认即开启第二步才能找到TFLite_Detection_PostProcess输出端口第二步转换时allow_custom_ops必须开启否则自定义后处理算子会导致转换失败若希望推理图与训练 / 评估配置在 NMS 阈值等参数上有所差异无需改动原配置直接通过--config_override覆盖即可。【免费下载链接】modelsModels and examples built with TensorFlow项目地址: https://gitcode.com/GitHub_Trending/mode/models创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考