如何在 Android 端优雅处理 TFLite 模型动态尺寸推理而无需频繁重构计算图

在移动端 AI 开发中，我们经常遇到输入尺寸不固定的场景，如 OCR 识别（文本行长度不一）、超分辨率（图片尺寸各异）或音频处理（时长不同）。传统的做法是根据最大尺寸进行 Padding，但这会造成严重的计算浪费；或者针对每个尺寸重新初始化解释器（Interpreter），但这会导致耗时的内存重新分配和模型解析。

本文将介绍如何利用 TensorFlow Lite 的 resizeInput 功能，在 Android 端实现高效的动态尺寸推理。

1. 准备支持动态 Shape 的模型

在导出模型时，我们需要将输入维度的对应位置设置为 None（即 -1）。

import tensorflow as tf

# 假设一个简单的卷积模型
model = tf.keras.Sequential([
    tf.keras.layers.InputLayer(input_shape=(None, None, 3), name='input'),
    tf.keras.layers.Conv2D(16, (3, 3), padding='same', activation='relu'),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 转化为 TFLite 格式
converter = tf.lite.TFLiteConverter.from_keras_model(model)
# 关键点：不要在转换时锁定 Shape
tflite_model = converter.convert()

with open('dynamic_model.tflite', 'wb') as f:
    f.write(tflite_model)

2. Android 端核心实现：ResizeInput

在 Android 端，我们不需要频繁销毁 Interpreter。当输入尺寸发生变化时，只需调用 resizeInput 方法，然后重新调用 allocateTensors。

Java/Kotlin 代码示例

// 1. 初始化 Interpreter
val options = Interpreter.Options()
val interpreter = Interpreter(loadModelFile(), options)

fun runInference(inputBitmap: Bitmap) {
    val width = inputBitmap.width
    val height = inputBitmap.height

    // 2. 动态调整输入 Tensor 尺寸
    // 假设输入索引为 0，Shape 为 [1, height, width, 3]
    val inputShape = intArrayOf(1, height, width, 3)
    interpreter.resizeInput(0, inputShape)

    // 3. 重新分配内存（这一步是必须的，否则会报错）
    // 注意：TFLite 会尽可能复用之前的内存，避免完全重新分配
    interpreter.allocateTensors()

    // 4. 准备输入数据
    val inputBuffer = convertBitmapToBuffer(inputBitmap, width, height)
    val outputMap = mutableMapOf<Int, Any>()
    val results = Array(1) { FloatArray(10) }
    outputMap[0] = results

    // 5. 执行推理
    interpreter.runForMultipleInputsOutputs(arrayOf(inputBuffer), outputMap)
}

3. 性能优化要点

1. 避免每帧都 Allocate：如果连续多帧的输入尺寸完全一致（例如视频流），则不需要再次调用 resizeInput 和 allocateTensors。维护一个 lastWidth 和 lastHeight 变量进行判断。
2. 使用 Hardware Buffer：在处理超大尺寸图片时，频繁的 allocateTensors 可能会触发内存碎片整理。建议配合 GPU Delegate 使用，GPU Delegate 在某些版本下对动态 Shape 的支持更好，能有效降低 CPU 负担。
3. 预设 Bucket：对于波动极小的输入，可以手动将其对齐到 8 或 16 的倍数（Padding），减少 resizeInput 的频率，从而在灵活性与性能间取得平衡。

总结

通过 resizeInput 和 allocateTensors 的配合，开发者可以优雅地在 Android 端处理变长输入，这比 Padding 到最大尺寸更节省算力，也比重新实例化解释器更流畅。这是端侧部署从“能跑”到“好用”的关键进阶技术。