使用tflite_runtime执行图像Captioning模型报错：不支持标准TensorFlow算子

环境与问题概述 #

• Windows 11系统，通过命令 !pip install --extra-index-url https://google-coral.github.io/py-repo/ tflite_runtime 安装了tflite_runtime 2.5.0.post1
• 运行自定义图像字幕模型时触发RuntimeError，但使用完整TensorFlow的tensorflow.lite模块可正常执行

推理代码 #

import numpy as np
from PIL import Image
import tflite_runtime.interpreter as tflite
from keras.preprocessing.sequence import pad_sequences  # 补充缺失的导入

max_len = 20
word_to_idx = np.load('weights/word_to_idx.npy', allow_pickle=True).item()
idx_to_word = np.load('weights/idx_to_word.npy', allow_pickle=True).item()

FEATURE_GENERATION_MODEL_TFLITE = 'feature_generation_model.tflite'
CAPTION_GENERATION_MODEL_TFLITE = 'caption_generation_model.tflite'

def predict_caption(path):
    a = Image.open(path)
    a = a.resize((300, 300))
    a = np.asarray(a, dtype='float32')
    imgp = a.reshape(1, 300, 300, 3)
    
    # 特征提取模型
    feat_interpreter = tflite.Interpreter(model_path=FEATURE_GENERATION_MODEL_TFLITE)
    feat_interpreter.allocate_tensors()

    input_index = feat_interpreter.get_input_details()[0]['index']
    output_index = feat_interpreter.get_output_details()[0]['index']

    feat_interpreter.set_tensor(input_index, imgp)
    feat_interpreter.invoke()
    
    feature_vector = feat_interpreter.get_tensor(output_index)
    feature_vector = feature_vector.reshape((1, 1536))
    
    # 生成字幕
    in_text = 'startseq'
    
    for i in range(max_len):
        seq = [word_to_idx[w] for w in in_text.split() if w in word_to_idx]
        seq = pad_sequences([seq], maxlen=max_len, padding='post')
        
        # 字幕生成模型
        cap_interpreter = tflite.Interpreter(model_path=CAPTION_GENERATION_MODEL_TFLITE)
        cap_interpreter.allocate_tensors()

        input_index1 = cap_interpreter.get_input_details()[0]['index']
        input_index2 = cap_interpreter.get_input_details()[1]['index']
        output_index = cap_interpreter.get_output_details()[0]['index']
        
        cap_interpreter.set_tensor(input_index1, feature_vector)
        cap_interpreter.set_tensor(input_index2, np.float32(seq))
        cap_interpreter.invoke()
        
        y_pred = cap_interpreter.get_tensor(output_index)
        y_pred = y_pred.argmax()
        
        word = idx_to_word[y_pred]
        in_text += ' '+word
        
        if word == 'endseq':
            break

    final_caption = in_text.split()[1:-1]
    final_caption = ' '.join(final_caption)

    return final_caption

报错信息 #

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
~\AppData\Local\Temp\ipykernel_10744\846162487.py in <module>
----> 1 predict_caption('images/image.jpg')

~\AppData\Local\Temp\ipykernel_10744\3775461012.py in predict_caption(path)
     91         cap_interpreter.set_tensor(input_index1, feature_vector)
     92         cap_interpreter.set_tensor(input_index2, np.float32(seq))
---> 93         cap_interpreter.invoke()
     94 
     95         y_pred = cap_interpreter.get_tensor(output_index)

~\anaconda3\lib\site-packages\tflite_runtime\interpreter.py in invoke(self)
    831     """
    832     self._ensure_safe()
---> 833     self._interpreter.Invoke()
    834 
    835   def reset_all_variables(self):

RuntimeError: Regular TensorFlow ops are not supported by this interpreter. Make sure you apply/link the Flex delegate before inference.Node number 9 (FlexTensorListReserve) failed to prepare.

TFLite模型转换代码 #

# 特征提取模型转换
FEATURE_GENERATION_MODEL_TFLITE = 'feature_generation_model.tflite'

tf_lite_converter = tf.lite.TFLiteConverter.from_keras_model(feature_generation_model)
feature_tflite_model = tf_lite_converter.convert()
open(FEATURE_GENERATION_MODEL_TFLITE, 'wb').write(feature_tflite_model)

# 字幕生成模型转换
CAPTION_GENERATION_MODEL_TFLITE = 'caption_generation_model.tflite'

tf_lite_converter = tf.lite.TFLiteConverter.from_keras_model(image_captioning_model)

tf_lite_converter.optimizations = [tf.lite.Optimize.DEFAULT]
tf_lite_converter.experimental_new_converter = True
tf_lite_converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]

caption_tflite_model = tf_lite_converter.convert()
open(CAPTION_GENERATION_MODEL_TFLITE, 'wb').write(caption_tflite_model)

解决方案 #

原因分析 #

字幕生成TFLite模型包含TFLite内置算子不支持的TensorFlow原生算子（如TensorListReserve），而tflite_runtime默认未集成Flex Delegate来兼容这类算子；完整TensorFlow自带Flex Delegate，因此可以正常运行。

解决步骤 #

1. 升级tflite_runtime版本
   2.5.0版本较旧，建议安装最新兼容版本：

pip install --extra-index-url https://google-coral.github.io/py-repo/ tflite_runtime

2. 加载模型时指定Flex Delegate
   修改字幕生成模型的初始化代码，添加Flex Delegate支持：

# 替换原字幕模型初始化代码
cap_interpreter = tflite.Interpreter(
    model_path=CAPTION_GENERATION_MODEL_TFLITE,
    experimental_delegates=[tflite.load_delegate('tensorflowlite_flex.dll')]
)

Windows系统下Flex Delegate动态库名为tensorflowlite_flex.dll，若找不到需确认tflite_runtime安装包包含该文件，或从TensorFlow官方下载对应版本的Flex Delegate库。

3. 优化模型转换（可选）
   检查字幕生成模型结构，替换无法被TFLite转换的算子；若模型可完全转换为TFLite内置算子，可移除tf.lite.OpsSet.SELECT_TF_OPS配置，避免依赖原生TF算子。

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