如何将fizyr Keras RetinaNet训练的h5模型转换为适配Android的TFLite模型

背景 #

使用fizyr开源Keras RetinaNet训练的类别识别模型，需集成到Android应用，初始采用如下脚本转换模型：

import tensorflow as tf
from keras_retinanet.models import load_model
from keras.layers import Input
from keras.models import Model
if __name__ == "__main__":
    model = load_model("modelFINAL.h5")
    fixed_input = Input((1080,1920,3))
    fixed_model = Model(fixed_input,model(fixed_input))
    converter = tf.lite.TFLiteConverter.from_keras_model(fixed_model)
    tflite = converter.convert()
    # 保存模型
    with open('model.tflite', 'wb') as f:
      f.write(tflite)

转换完成后导入Android应用触发崩溃，核心报错信息为：

Didn't find op for builtin opcode 'MUL' version '5'. An older version of this builtin might be supported. Are you using an old TFLite binary with a newer model?

错误根源 #

报错核心为TFLite模型算子版本和Android端运行库版本不匹配：本地转换模型使用的TensorFlow版本偏高，生成的TFLite模型用到了v5版本的MUL算子，但Android项目集成的TFLite运行库版本较低，无法识别高版本算子。

可落地解决方案 #

方案1：调整转换参数适配低版本运行库（无需修改Android端代码，优先推荐） #

在原有转换脚本中添加算子版本控制配置，强制转换器生成低版本运行库可兼容的算子，修改后代码如下：

import tensorflow as tf
from keras_retinanet.models import load_model
from keras.layers import Input
from keras.models import Model

if __name__ == "__main__":
    # 加载模型时需传入训练时对应的backbone名称，默认训练用的是resnet50
    model = load_model("modelFINAL.h5", backbone_name='resnet50')
    fixed_input = Input((1080,1920,3))
    fixed_model = Model(fixed_input, model(fixed_input))
    
    converter = tf.lite.TFLiteConverter.from_keras_model(fixed_model)
    # 指定仅使用TFLite内置算子，禁用高版本独有算子
    converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
    converter._experimental_lower_tensor_list_ops = False
    # 若仍有算子不兼容，可开启下一行使用TF算子兼容模式
    # converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS)
    
    tflite = converter.convert()
    # 保存转换后模型
    with open('model.tflite', 'wb') as f:
        f.write(tflite)

如果开启了SELECT_TF_OPS配置，需要在Android端的build.gradle中添加对应依赖：

implementation 'org.tensorflow:tensorflow-lite-select-tf-ops:2.8.0'

方案2：升级Android端TFLite依赖版本 #

将Android项目build.gradle中的TFLite相关依赖升级到和本地转换用的TensorFlow大版本完全一致即可，比如本地用的是TensorFlow 2.8版本，就修改为：

implementation 'org.tensorflow:tensorflow-lite:2.8.0'
implementation 'org.tensorflow:tensorflow-lite-task-vision:2.8.0'

注意事项 #

• RetinaNet模型加载时必须指定backbone_name参数，和训练时使用的骨干网络一致，否则会出现模型加载错误
• 优先选择方案1，不会增大安装包体积，适配性更强
• 转换前可以先执行print(tf.__version__)查看本地TensorFlow版本，确保和Android端版本匹配

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