将TF Object Detection API模型转TFLite时遇google3模块缺失错误求助

问题复现 #

执行以下导出脚本时触发错误：

!python models/research/object_detection/export_tflite_graph_tf2.py \
    --trained_checkpoint_dir {'/content/drive/MyDrive/customTF2/training'} \
    --output_directory {'/content/drive/MyDrive/customTF2/tflite'} \
    --pipeline_config_path {'/content/drive/MyDrive/customTF2/data/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.config'}

错误信息：

File "models/research/object_detection/export_tflite_graph_tf2.py", line 99, in <module>
    from google3.image.understanding.object_detection.input_readers import input_reader_google_pb2   #pylint: disable=unused-import
ModuleNotFoundError: No module named 'google3'

修复方案 #

方案1：直接修改导出脚本 #

找到export_tflite_graph_tf2.py文件的第99行，注释掉导入google3的语句——这个模块是Google内部专用组件，对外用户完全不需要。修改后该行变为：

# from google3.image.understanding.object_detection.input_readers import input_reader_google_pb2   #pylint: disable=unused-import

方案2：使用官方推荐的替代导出方式 #

Google官方目前更推荐通过SavedModel中转的方式转换TFLite，步骤如下：

1. 加载训练模型与配置

import tensorflow as tf
from object_detection.utils import config_util
from object_detection.builders import model_builder

# 加载配置文件
configs = config_util.get_configs_from_pipeline_file('/content/drive/MyDrive/customTF2/data/ssd_mobilenet_v2_fpnlite_320x320_coco17_tpu-8.config')
model_config = configs['model']
detection_model = model_builder.build(model_config=model_config, is_training=False)

# 加载训练好的checkpoint（替换XXXX为实际的checkpoint编号）
ckpt = tf.compat.v2.train.Checkpoint(model=detection_model)
ckpt.restore('/content/drive/MyDrive/customTF2/training/ckpt-XXXX').expect_partial()

# 保存为SavedModel格式
tf.saved_model.save(detection_model, '/content/drive/MyDrive/customTF2/saved_model')

2. 转换为TFLite模型

# 初始化转换器
converter = tf.lite.TFLiteConverter.from_saved_model('/content/drive/MyDrive/customTF2/saved_model')
# 开启默认优化
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 支持TF自定义算子
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]

# 执行转换
tflite_model = converter.convert()

# 保存最终的TFLite模型
with open('/content/drive/MyDrive/customTF2/tflite/detect.tflite', 'wb') as f:
    f.write(tflite_model)

方案3：更新TF Object Detection API版本 #

旧版本的TFOD API可能残留了内部模块引用，执行以下命令更新到最新版本：

cd models/research/
git pull

内容的提问来源于stack exchange，提问作者theekshana1987