TF-Vision模型转TFLite后输入尺寸异常的解决方法咨询

从TensorFlow Model Garden训练完成的RetinaNet-Mobile目标检测模型，使用export_tflite_lib.convert_tflite_model转换为TFLite格式后，模型输入尺寸显示为[1,1,1,3]，无法适配256x256的输入图像，预期输入尺寸应为[1,None,None,3]（动态尺寸）或固定[1,256,256,3]。

环境配置 #

pip install tensorflow==2.13
pip install imagesize opencv-python chardet
pip install tf-models-official # the tf-vision libs

复现代码 #

import tensorflow as tf
import tensorflow_models as tfm
from official.vision.serving import export_saved_model_lib, export_tflite_lib
import PIL.Image
import numpy as np
import os, time, json, imagesize, sys

print("TF", tf.__version__) # 2.13

MODEL_WIDTH = 256
MODEL_HEIGHT = 256
BATCH_SIZE = 1

# 训练完成模型的导出文件夹
EXPORT_DIR = "/home/.../202309_tfvision/export"

# ------------------------------------------------------------------------------------
# 加载实验配置
exp_config = tfm.core.exp_factory.get_exp_config('retinanet_mobile_coco')

logical_device_names = [logical_device.name for logical_device in tf.config.list_logical_devices()]

# 骨干网络配置
exp_config.task.freeze_backbone = False
exp_config.task.annotation_file = ''

# 模型配置
exp_config.task.model.input_size = [MODEL_WIDTH, MODEL_HEIGHT, 3]
exp_config.task.model.num_classes = 3
exp_config.task.model.detection_generator.tflite_post_processing.max_classes_per_detection = exp_config.task.model.num_classes

train_steps = 1 # 测试用

# 训练器配置调整
if 'GPU' in ''.join(logical_device_names):
  print('This may be broken in Colab.')
  device = 'GPU'
elif 'TPU' in ''.join(logical_device_names):
  print('This may be broken in Colab.')
  device = 'TPU'
else:
  print('Running on CPU is slow, so only train for a few steps.')
  device = 'CPU'

exp_config.trainer.steps_per_loop = 10
exp_config.trainer.summary_interval = 100
exp_config.trainer.checkpoint_interval = 4
exp_config.trainer.validation_interval = 8
exp_config.trainer.validation_steps =  1
exp_config.trainer.train_steps = train_steps

exp_config.trainer.optimizer_config.warmup.linear.warmup_steps = 100
exp_config.trainer.optimizer_config.learning_rate.type = 'cosine'
exp_config.trainer.optimizer_config.learning_rate.cosine.decay_steps = train_steps
exp_config.trainer.optimizer_config.learning_rate.cosine.initial_learning_rate = 0.1
exp_config.trainer.optimizer_config.warmup.linear.warmup_learning_rate = 0.05

print(exp_config.as_dict())

# ------------------------------------------------------------------------------------
print("TFLITE")

tflite_model = export_tflite_lib.convert_tflite_model(
    saved_model_dir=EXPORT_DIR,
    params=exp_config,
    quant_type=None)

interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

print("input_details", input_details)

# 测试TFLite模型
TEST_IMAGE_FILE = "/home/daniel/armazem/porthus/ml_data/camera/20230731_233104/poco/000165.jpg"

image = np.array(PIL.Image.open(TEST_IMAGE_FILE).resize((MODEL_WIDTH, MODEL_HEIGHT)), dtype=input_details[0]["dtype"])
images= np.expand_dims(image, axis=0)

# 报错：ValueError: Cannot set tensor: Dimension mismatch. Got 256 but expected 1 for dimension 1 of input 0.
interpreter.set_tensor(input_details[0]["index"], images)
interpreter.invoke()

解决方案 #

1. 显式指定TFLite转换的输入形状 #

在调用convert_tflite_model时添加input_shape参数，明确固定输入尺寸：

tflite_model = export_tflite_lib.convert_tflite_model(
    saved_model_dir=EXPORT_DIR,
    params=exp_config,
    quant_type=None,
    input_shape=[BATCH_SIZE, MODEL_WIDTH, MODEL_HEIGHT, 3]
)

2. 修正实验配置的输入尺寸层级 #

部分Model Garden模型的输入尺寸配置需要设置在task层级而非task.model，尝试调整配置：

# 替换原来的exp_config.task.model.input_size设置
exp_config.task.input_size = [MODEL_HEIGHT, MODEL_WIDTH, 3]

3. 验证SavedModel的输入签名（排查根源） #

先确认导出的SavedModel本身输入形状是否正确，避免转换TFLite前就存在问题：

loaded_model = tf.saved_model.load(EXPORT_DIR)
infer_signature = loaded_model.signatures['serving_default']
print("SavedModel input shape:", infer_signature.inputs[0].shape)

如果SavedModel输入就是[1,1,1,3]，需要重新导出SavedModel时确保input_size配置被正确应用。

4. 临时适配：手动调整TFLite输入尺寸 #

如果上述方法无法立即生效，可在加载TFLite模型后手动修改输入尺寸作为临时方案：

interpreter = tf.lite.Interpreter(model_content=tflite_model)
# 调整输入张量尺寸
interpreter.resize_tensor_input(input_details[0]['index'], [1, MODEL_WIDTH, MODEL_HEIGHT, 3])
interpreter.allocate_tensors()

# 后续测试代码不变
image = np.array(PIL.Image.open(TEST_IMAGE_FILE).resize((MODEL_WIDTH, MODEL_HEIGHT)), dtype=input_details[0]["dtype"])
images= np.expand_dims(image, axis=0)
interpreter.set_tensor(input_details[0]["index"], images)
interpreter.invoke()

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