在Android Studio部署float32类型TFlite模型如何定义推理用ByteBuffer

Float32类型TFLite模型ByteBuffer实现方案 #

Float32模型的ByteBuffer不需要Int8量化相关的缩放、零点参数，只需按照输入张量的元素总数、每个float占4字节的规则分配容量，写入归一化后的像素值即可，具体实现如下：

核心逻辑说明 #

• 容量计算规则：输入张量总元素数 * 4字节，你当前使用的输入形状为[1,256,256,3]，所以容量为1 * 256 * 256 * 3 * 4 = 786432字节
• 必须配置nativeOrder()字节序，匹配TFLite底层运算的字节要求
• 归一化逻辑必须和模型训练阶段的预处理完全一致，否则会出现推理结果异常

完整代码实现（Kotlin，以图片输入场景为例） #

import android.graphics.Bitmap
import android.graphics.Color
import java.nio.ByteBuffer
import java.nio.ByteOrder

// 1. 先将原始输入Bitmap缩放到模型要求的256x256尺寸
val scaledBitmap = Bitmap.createScaledBitmap(originalBitmap, 256, 256, true)

// 2. 初始化ByteBuffer
val byteBuffer = ByteBuffer.allocateDirect(1 * 256 * 256 * 3 * 4)
byteBuffer.order(ByteOrder.nativeOrder())

// 3. 提取Bitmap所有像素值到Int数组
val pixelArray = IntArray(256 * 256)
scaledBitmap.getPixels(pixelArray, 0, scaledBitmap.width, 0, 0, scaledBitmap.width, scaledBitmap.height)

var index = 0
// 4. 遍历像素写入Float值到ByteBuffer
for (h in 0 until 256) {
    for (w in 0 until 256) {
        val pixel = pixelArray[index++]
        // 提取ARGB格式下的R、G、B三通道值
        val r = Color.red(pixel)
        val g = Color.green(pixel)
        val b = Color.blue(pixel)

        // 【此处必须和训练预处理对齐】下面提供两种常用归一化写法，按需选择/修改
        // 写法1：训练时将像素除以255归一化到[0,1]区间
        byteBuffer.putFloat(r / 255f)
        byteBuffer.putFloat(g / 255f)
        byteBuffer.putFloat(b / 255f)

        // 写法2：训练时将像素归一化到[-1,1]区间，替换上面三行即可
        // byteBuffer.putFloat((r / 127.5f) - 1f)
        // byteBuffer.putFloat((g / 127.5f) - 1f)
        // byteBuffer.putFloat((b / 127.5f) - 1f)
    }
}

// 5. 重置ByteBuffer游标到起始位置，否则TFLite无法正常读取
byteBuffer.rewind()

// 后续直接调用原有代码的inputFeature0.loadBuffer(byteBuffer)即可

注意事项 #

• 通道顺序对齐：如果模型训练时使用的是BGR通道输入，将R、G、B的写入顺序反过来即可
• 输入尺寸对齐：如果原始输入不是256x256，除了缩放外必要时需要做中心裁剪等处理，保证输入形状完全匹配模型要求
• 自定义预处理兼容：如果训练时使用了均值、方差做标准化，修改putFloat内的计算逻辑即可，比如(r - 均值R) / 方差R

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