Android端TensorFlow Lite模型整数数组输入的推理实现方法

当前开发场景说明 #

• 已通过Android Studio的TFLite模型导入流程完成模型集成，导入操作路径为(New -> Other -> Tensorflow Lite Model)，可直接使用自动生成的模型封装类开展开发
• 模型核心规格：
  • 输入：推理时批大小固定为1，实际输入维度1×100×12，数据类型为整数
  • 输出：推理时批大小固定为1，实际输出维度1×3，数据类型为32位浮点
• 当前卡点：现有自动生成的示例代码默认适配图像类的浮点输入，既无法通过loadArray接口直接接收int[1][100][12]三维数组，也不清楚loadBuffer接口的正确传参方式

正确实现步骤 #

1. 修正TensorBuffer的创建参数 #

原示例代码中创建TensorBuffer时默认指定为DataType.FLOAT32，与模型的整数输入类型不匹配，需要替换为和模型输入完全一致的DataType：

• 如果导出TFLite模型时输入为32位整数，对应使用DataType.INT32
• 如果是量化后的8位无符号整数输入，对应使用DataType.UINT8
  错误示例的修正前后对比：

// 错误：默认浮点类型，和模型整数输入不匹配
TensorBuffer inputFeature0 = TensorBuffer.createFixedSize(new int[]{1, 100, 12}, DataType.FLOAT32);

// 正确：根据模型实际输入类型调整，以INT32整数输入为例
TensorBuffer inputFeature0 = TensorBuffer.createFixedSize(new int[]{1, 100, 12}, DataType.INT32);

2. 使用loadArray接口的正确写法 #

TensorBuffer的loadArray接口不支持直接传入三维数组，只支持传入一维的基本类型数组，需要把三维的原始输入按维度顺序展平为一维数组即可，总元素数为1*100*12=1200，不需要手动计算内存偏移。
示例代码：

// 原始采集的三维输入数据，例如人体活动识别的传感器序列
int[][][] rawSensorData = ...;

// 按原维度顺序展平为一维int数组
int[] flatInput = new int[1 * 100 * 12];
int writeOffset = 0;
for (int i = 0; i < rawSensorData.length; i++) {
    for (int j = 0; j < rawSensorData[i].length; j++) {
        int[] currentRow = rawSensorData[i][j];
        System.arraycopy(currentRow, 0, flatInput, writeOffset, currentRow.length);
        writeOffset += currentRow.length;
    }
}

// 直接加载展平后的数组
inputFeature0.loadArray(flatInput);

3. 使用loadBuffer接口的正确写法 #

如果选择手动构造ByteBuffer传入，需要严格遵守以下规则避免运行时错误：

• ByteBuffer需要调用allocateDirect创建，总容量等于 总元素数 * 单元素字节长度：例如INT32类型每个元素占4字节，总容量为1200 * 4 = 4800；UINT8类型每个元素占1字节，总容量为1200 * 1 = 1200
• 必须手动设置字节序为ByteOrder.nativeOrder()，和设备的原生字节序保持一致，避免大小端错乱导致的值读取错误
• 写入完所有数据后必须调用rewind()方法把缓冲区指针移到起始位置，否则loadBuffer会从当前指针位置开始读取，出现数据越界
  示例代码：

// INT32类型的ByteBuffer构造示例
ByteBuffer inputBuffer = ByteBuffer.allocateDirect(1 * 100 * 12 * 4);
inputBuffer.order(ByteOrder.nativeOrder());

// 按原始维度顺序逐值写入
for (int[][][] rawData ...) {
    for (int[] row : rawData) {
        for (int val : row) {
            inputBuffer.putInt(val);
        }
    }
}
inputBuffer.rewind();
inputFeature0.loadBuffer(inputBuffer);

常见避坑点 #

• 展平数组时必须严格对齐训练时的输入顺序：即按批维度、时间步维度、特征维度的顺序拼接，不要随意调整行或列的顺序，否则推理结果会完全偏离预期
• 不要随意修改输入的维度顺序，哪怕维度总数一致，顺序不对也会触发模型运行时的形状校验错误
• 模型使用完成后依然要按照原逻辑调用model.close()释放推理资源，避免内存泄漏

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