2. TensorFlow Lite模型量化与部署到Android

任务目标

本案例将完成一个完整的端到端流程：训练一个简单的图像分类模型（基于MNIST数据集），将其转换为TensorFlow Lite格式并应用动态范围量化，最终在Android应用中加载该模型进行手写数字识别。通过本案例，你将掌握TFLite模型的生成、量化及Android端部署的核心步骤。

环境与依赖

Python端：TensorFlow 2.x（建议2.13+）、numpy
Android端：Android Studio（Hedgehog或更新版本）、minSdk 21+、TFLite Runtime 2.14+

第一步：训练模型并转换为量化TFLite

以下Python脚本完成模型训练、动态范围量化转换，并保存为mnist_quant.tflite。动态范围量化仅将权重从float32转为int8，激活值仍为float，在保持较高精度的同时将模型体积缩小约4倍。

import tensorflow as tf
import numpy as np

# 1. 加载MNIST数据
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train = x_train.reshape(-1, 28, 28, 1).astype('float32') / 255.0
x_test = x_test.reshape(-1, 28, 28, 1).astype('float32') / 255.0

# 2. 构建并训练模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(8, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(x_train, y_train, epochs=2, batch_size=32, validation_split=0.1, verbose=1)

# 3. 转换为TFLite并应用动态范围量化
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]  # 动态范围量化
tflite_model = converter.convert()

# 4. 保存模型
with open('mnist_quant.tflite', 'wb') as f:
    f.write(tflite_model)

# 5. 验证转换后的模型（可选）
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 用测试集第一张图验证
test_input = x_test[0:1]
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
pred = interpreter.get_tensor(output_details[0]['index'])
print(f'预期标签: {y_test[0]}, 预测结果: {np.argmax(pred)}')
print(f'模型大小: {len(tflite_model)} bytes')

预期输出：训练2个epoch后准确率约97%以上；验证代码输出预期标签: 7, 预测结果: 7（MNIST测试集第一张图标签为7）；模型大小约20KB（原始float32模型约80KB）。

第二步：Android项目集成

在Android Studio中创建新项目，选择Empty Views Activity，包名设为com.example.tflite_mnist。将生成的mnist_quant.tflite文件放入app/src/main/assets/目录。

在app/build.gradle中添加依赖：

dependencies {
    implementation 'org.tensorflow:tensorflow-lite:2.14.0'
    // 如需GPU加速可添加: implementation 'org.tensorflow:tensorflow-lite-gpu:2.14.0'
}

第三步：Android推理代码

创建一个工具类Classifier.kt，负责加载模型和执行推理。以下代码实现了从Bitmap输入到分类结果的完整流程，输入为28×28灰度图像，输出为10个类别的概率分布。

import android.content.Context
import android.graphics.Bitmap
import org.tensorflow.lite.Interpreter
import java.io.FileInputStream
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.nio.channels.FileChannel

class Classifier(context: Context) {
    private val interpreter: Interpreter

    init {
        // 从assets加载模型
        val modelBuffer = loadModelFile(context, "mnist_quant.tflite")
        interpreter = Interpreter(modelBuffer)
    }

    private fun loadModelFile(context: Context, modelPath: String): ByteBuffer {
        val assetFileDescriptor = context.assets.openFd(modelPath)
        val inputStream = FileInputStream(assetFileDescriptor.fileDescriptor)
        val fileChannel = inputStream.channel
        val startOffset = assetFileDescriptor.startOffset
        val declaredLength = assetFileDescriptor.declaredLength
        return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength)
    }

    // 将Bitmap转换为模型输入格式 (1, 28, 28, 1) float32
    private fun convertBitmapToInput(bitmap: Bitmap): ByteBuffer {
        val inputSize = 28 * 28 * 4  // float32占4字节
        val inputBuffer = ByteBuffer.allocateDirect(inputSize)
        inputBuffer.order(ByteOrder.nativeOrder())

        val scaledBitmap = Bitmap.createScaledBitmap(bitmap, 28, 28, true)
        val pixels = IntArray(28 * 28)
        scaledBitmap.getPixels(pixels, 0, 28, 0, 0, 28, 28)

        for (pixel in pixels) {
            // 将ARGB像素转为灰度值并归一化到[0,1]
            val r = (pixel shr 16) and 0xFF
            val g = (pixel shr 8) and 0xFF
            val b = pixel and 0xFF
            val gray = (0.299f * r + 0.587f * g + 0.114f * b) / 255.0f
            inputBuffer.putFloat(gray)
        }
        return inputBuffer
    }

    // 执行推理，返回10个类别的概率
    fun recognize(bitmap: Bitmap): FloatArray {
        val inputBuffer = convertBitmapToInput(bitmap)
        val outputBuffer = Array(1) { FloatArray(10) }

        interpreter.run(inputBuffer, outputBuffer)
        return outputBuffer[0]
    }

    fun close() {
        interpreter.close()
    }
}

第四步：MainActivity调用

在MainActivity.kt中，从资源或相册获取手写数字图片，调用Classifier进行识别并显示结果。以下代码演示了从drawable资源加载图片并推理的完整流程：

import android.graphics.BitmapFactory
import android.os.Bundle
import android.widget.Button
import android.widget.ImageView
import android.widget.TextView
import androidx.appcompat.app.AppCompatActivity

class MainActivity : AppCompatActivity() {
    private lateinit var classifier: Classifier

    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        setContentView(R.layout.activity_main)

        classifier = Classifier(this)

        val imageView = findViewById<ImageView>(R.id.imageView)
        val resultText = findViewById<TextView>(R.id.resultText)
        val recognizeBtn = findViewById<Button>(R.id.recognizeBtn)

        // 从drawable加载测试图片（需准备一张手写数字图片）
        val bitmap = BitmapFactory.decodeResource(resources, R.drawable.test_digit)
        imageView.setImageBitmap(bitmap)

        recognizeBtn.setOnClickListener {
            val results = classifier.recognize(bitmap)
            val maxIndex = results.indices.maxByOrNull { results[it] } ?: -1
            val confidence = results[maxIndex] * 100
            resultText.text = "识别结果: $maxIndex (置信度: ${"%.1f".format(confidence)}%)"
        }
    }

    override fun onDestroy() {
        super.onDestroy()
        classifier.close()
    }
}

对应的activity_main.xml布局包含一个ImageView、一个TextView和一个Button，ID分别为imageView、resultText、recognizeBtn。在res/drawable中放置一张28×28或更大尺寸的手写数字图片（如命名为test_digit.png），建议使用白底黑字。

验证方法与通过标准

Python端验证：运行转换脚本，观察输出中预测结果与预期标签一致，模型文件成功生成。
Android端验证：连接Android设备（或使用模拟器），点击Run安装应用。应用启动后显示测试图片，点击识别按钮，TextView显示识别结果。
通过标准：应用能正确识别测试图片中的手写数字（如测试图片为数字"7"，则显示"识别结果: 7"），且无崩溃或ANR。若识别错误，可检查图片预处理是否与训练时一致（灰度、归一化、尺寸）。

常见问题

模型加载失败：确认mnist_quant.tflite已放入assets目录且文件名拼写正确。
输入尺寸不匹配：模型输入为(1, 28, 28, 1)，确保预处理时Bitmap已缩放到28×28且转换为灰度float数组。
识别准确率低：检查图片是否为白底黑字（MNIST训练数据为黑底白字，但灰度反转后模型仍可识别，若效果差可尝试将像素值取反：gray = 1.0f - gray）。

📚 课程目录