Android加载GPT-2 TFLite模型时遇缓冲区大小不匹配错误

错误信息 #

java.lang.IllegalArgumentException: 无法从大小为393216字节的TensorFlowLite张量(StatefulPartitionedCall:8)复制到大小为512字节的Java Buffer中。

输入模型详情 #

[{'name': 'serving_default_input_1:0', 'index': 0, 'shape': array([ 1, 64], dtype=int32), 'shape_signature': array([ 1, 64], dtype=int32), 'dtype': <class 'numpy.int32'>, 'quantization': (0.0, 0), 'quantization_parameters': {'scales': array([], dtype=float32), 'zero_points': array([], dtype=int32), 'quantized_dimension': 0}, 'sparsity_parameters': {}}]

问题分析 #

1. 缓冲区大小不匹配：代码中硬编码输出缓冲区为128 * 4 = 512字节，但模型实际输出张量大小为393216字节，两者差距过大导致报错。
2. 文本生成逻辑错误：GPT-2是自回归模型，无法通过单次输入直接得到完整生成文本，需要循环预测下一个token并迭代更新输入序列。

解决方案 #

步骤1：获取模型输出张量的准确信息 #

在初始化模型时，添加代码打印输出张量的形状、数据类型和字节大小，避免硬编码缓冲区：

class LoadModel(context: Context) {
    private var tflite: Interpreter? = null
    private var outputTensorShape: IntArray? = null
    private var outputTensorByteSize: Int = 0

    init {
        try {
            tflite = Interpreter(loadModelFile(context, "gpt2_model.tflite"))
            // 获取输出张量信息
            tflite?.getOutputTensor(0)?.let { tensor ->
                outputTensorShape = tensor.shape()
                outputTensorByteSize = tensor.numBytes()
                println("Output tensor shape: ${outputTensorShape?.contentToString()}")
                println("Output tensor byte size: $outputTensorByteSize")
                println("Output tensor dtype: ${tensor.dataType()}")
            }
        } catch (e: Exception) {
            e.printStackTrace()
        }
    }

    // 原有loadModelFile方法保留
    @Throws(Exception::class)
    private fun loadModelFile(context: Context, modelName: String): MappedByteBuffer {
        val fileInputStream = FileInputStream(context.assets.openFd(modelName).fileDescriptor)
        val fileChannel = fileInputStream.channel
        val startOffset = context.assets.openFd(modelName).startOffset
        val declaredLength = context.assets.openFd(modelName).declaredLength
        return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength)
    }

    // 其他方法后续修改...
}

步骤2：修正输出缓冲区分配与生成逻辑 #

根据获取到的输出张量字节数动态分配缓冲区，并实现自回归文本生成逻辑：

fun generateText(inputText: String, maxLength: Int = 128): String {
    tflite ?: return "Model not initialized"
    outputTensorShape ?: return "Output tensor info not available"

    // 初始化输入序列
    var tokens = tokenize(inputText).toMutableList()
    val maxSeqLength = 64 // 匹配模型输入形状的序列长度

    // 截断或填充初始输入到指定长度
    if (tokens.size > maxSeqLength) {
        tokens = tokens.subList(tokens.size - maxSeqLength, tokens.size)
    } else {
        while (tokens.size < maxSeqLength) {
            tokens.add(0) // 用padding token填充
        }
    }

    // 生成文本循环
    while (tokens.size < maxLength) {
        // 准备输入缓冲区
        val inputBuffer = ByteBuffer.allocateDirect(maxSeqLength * 4)
        inputBuffer.order(ByteOrder.nativeOrder())
        tokens.takeLast(maxSeqLength).forEach { token ->
            inputBuffer.putInt(token)
        }
        inputBuffer.rewind()

        // 准备输出缓冲区（使用动态获取的字节大小）
        val outputBuffer = ByteBuffer.allocateDirect(outputTensorByteSize)
        outputBuffer.order(ByteOrder.nativeOrder())

        // 运行推理
        tflite!!.run(inputBuffer, outputBuffer)
        outputBuffer.rewind()

        // 解析输出，获取最后一个位置的logits（假设输出形状为[1, seq_len, vocab_size]）
        val vocabSize = outputTensorShape!![2]
        val logits = FloatArray(vocabSize)
        // 定位到最后一个token的logits位置
        outputBuffer.position((maxSeqLength - 1) * vocabSize * 4)
        for (i in 0 until vocabSize) {
            logits[i] = outputBuffer.getFloat()
        }

        // 采样得到下一个token（简单取概率最大的token）
        val nextToken = logits.indices.maxByOrNull { logits[it] } ?: 0
        if (nextToken == 50256) { // GPT-2的停止token
            break
        }

        // 添加新token到序列
        tokens.add(nextToken)
    }

    // 解码生成的token序列
    return detokenize(tokens.toIntArray())
}

// 原有close方法保留
fun close() {
    if (tflite != null) {
        tflite!!.close()
        tflite = null
    }
}

步骤3：实现正确的Tokenizer/Detokenizer #

必须实现与GPT-2匹配的tokenizer逻辑，示例框架如下：

private fun tokenize(inputText: String): IntArray {
    // 实现GPT-2的tokenize逻辑
    // 可使用Hugging Face Transformers for Android库，或手动加载vocab.json/merges.txt实现ByteLevelBPETokenizer
    return intArrayOf() // 替换为实际实现
}

private fun detokenize(tokens: IntArray): String {
    // 实现GPT-2的detokenize逻辑
    return "Generated text" // 替换为实际实现
}

注意事项 #

• 模型转换时需确保保留LM Head，否则输出会是隐藏层特征而非logits，无法直接用于文本生成
• 自回归生成会多次调用模型，可启用TFLite GPU delegate提升性能
• 可优化采样策略（如top-k、top-p采样），避免生成重复文本

内容的提问来源于stack exchange，提问作者Subha lakshmi S