Android TensorFlow Lite(nightly)下[H,W,C]转[B,C,H,W]实现方法问询

文档查询路径 #

TensorFlow Lite 提供的Java/Kotlin层API的所有张量操作相关定义，都可以通过两个途径查询：

• 引入tensorflow-lite和tensorflow-lite-support依赖后，直接在IDE中查看对应类的官方注释文档
• 下载TensorFlow Lite Android SDK离线包后，在解压目录的docs/api_reference/java子目录下查看完整的离线API参考

对应维度操作的实现参考 #

你需要实现的[H,W,C]转[B,C,H,W]的需求，可以通过TensorFlow Lite Support Library提供的TensorProcessor工具类完成，和你之前C++端的操作逻辑对应，示例代码如下：

import org.tensorflow.lite.support.common.ops.ExpandDimsOp
import org.tensorflow.lite.support.image.ops.TransposeOp
import org.tensorflow.lite.support.tensorbuffer.TensorBuffer
import org.tensorflow.lite.support.common.TensorProcessor

// 假设inputTensorBuffer为你已获取的形状[H,W,C]的输入张量
// 第一步：执行维度置换，对应C++端permute({2,0,1})，输出形状为[C,H,W]
val transposedTensor = TensorProcessor.Builder()
    .add(TransposeOp(intArrayOf(2, 0, 1)))
    .build()
    .process(inputTensorBuffer)

// 第二步：在0号位置添加batch维度，对应C++端unsqueeze(0)，输出形状为[1,C,H,W]即[B,C,H,W]
val resultTensor = TensorProcessor.Builder()
    .add(ExpandDimsOp(0))
    .build()
    .process(transposedTensor)

常用维度操作映射 #

你提到的几个操作在Java/Kotlin层的对应实现如下：

• transpose/permute：对应TransposeOp，构造参数为维度置换顺序的int数组
• unsqueeze：对应ExpandDimsOp，构造参数为要插入维度的索引位
• squeeze：对应SqueezeOp，可传入指定要移除的大小为1的维度索引，不传则默认移除所有大小为1的维度

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