如何将seq2seq注意力模型转换为可在Android Studio部署的.tflite格式？

问题原因说明 #

• 导入Android Studio提示「不是有效的tflite模型文件」主要有三个原因：
  1. 转换过程未固定输入输出的序列长度，动态shape无法被Android Studio的模型校验工具识别
  2. 模型使用了TensorFlow原生算子（SELECT_TF_OPS），但Android端未添加对应支持依赖
  3. 打包时tflite文件被默认压缩，导致模型结构被破坏
• 模型体积接近200MB上限的核心原因是导出时把分词词表、前后处理逻辑都内嵌到了SavedModel中，这部分冗余内容占了模型体积的大部分。

具体落地步骤 #

1. 拆分前后处理逻辑 #

将原模型中的分词、词表映射逻辑从SavedModel中剥离，放到Android端用原生代码实现，导出的模型仅保留Encoder、Decoder注意力推理的核心部分，可直接砍掉70%以上的冗余体积。

2. 优化TFLite转换流程 #

参考如下TF2.X代码完成转换，固定输入输出shape并开启全整数量化：

# 固定输入序列长度为业务场景下的最大支持长度，示例为64个token
input_tokens = tf.keras.Input(shape=(64,), dtype=tf.int32, name="input_tokens")

# 封装推理逻辑，固定输出shape，避免动态维度
@tf.function(input_signature=[input_tokens])
def infer_pipeline(input_tokens):
    translate_result = translator.tf_translate(input_tokens)
    # 固定输出序列长度，和输入长度对齐
    return tf.ensure_shape(translate_result['text'], (None, 64))

# 保存仅含核心推理逻辑的SavedModel
tf.saved_model.save(
    translator,
    'lite_saved_model',
    signatures={'serving_default': infer_pipeline}
)

# 初始化转换器
converter = tf.lite.TFLiteConverter.from_saved_model('lite_saved_model')
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS,
  tf.lite.OpsSet.SELECT_TF_OPS
]

# 开启全整数量化，大幅压缩体积
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# 准备校准数据集，取100-200条训练样本即可
def representative_data_gen():
    for sample in calibration_dataset.take(100):
        yield [sample]
converter.representative_dataset = representative_data_gen
converter.target_spec.supported_types = [tf.int8]
converter.experimental_enable_resource_variables = True
converter.experimental_allow_tf_ops_fallback = True

# 转换并保存模型
tflite_model = converter.convert()
with open('translator.tflite', 'wb') as f:
    f.write(tflite_model)

3. Android端适配 #

在模块级build.gradle中添加如下配置，避免压缩并添加自定义算子支持：

android {
    aaptOptions {
        noCompress "tflite"
    }
}

dependencies {
    implementation 'org.tensorflow:tensorflow-lite:2.15.0'
    implementation 'org.tensorflow:tensorflow-lite-select-tf-ops:2.15.0'
    // 如需GPU加速可添加以下依赖
    // implementation 'org.tensorflow:tensorflow-lite-gpu:2.15.0'
}

按以上流程处理后的TFLite模型体积一般可压缩到50MB以内，可直接导入Android Studio正常部署使用。

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