从零开始：使用cnfood-241数据集训练基础模型并导出int8量化.tflite文件实战指南

1次阅读

没有评论

共计 2677 个字符，预计需要花费 7 分钟才能阅读完成。

在移动端部署机器学习模型时，我们常常面临两个主要挑战：模型体积过大和推理速度慢。大模型会占用宝贵的存储空间，增加应用安装包大小，而缓慢的推理速度则直接影响用户体验。这些问题在资源受限的移动设备上尤为突出。

量化是解决上述问题的有效方法之一。常见的量化方式包括：

float32 到 float16 的量化：能减少模型体积，但对推理速度提升有限
动态范围量化：简单易用，但优化效果一般
int8 量化：显著减小模型体积（通常减少 75%），同时大幅提升推理速度（通常 2 - 3 倍）

int8 量化通过将模型参数从 32 位浮点数转换为 8 位整数，在保持较高精度的同时获得最佳的性能优化效果，因此成为我们的首选方案。

首先我们需要准备 cnfood-241 数据集。这里提供一个完整的预处理示例：

import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator

# 数据路径
train_dir = 'path_to_cnfood241/train'
val_dir = 'path_to_cnfood241/val'

# 数据增强配置
train_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True,
    validation_split=0.2
)

# 训练数据生成器
train_generator = train_datagen.flow_from_directory(
    train_dir,
    target_size=(224, 224),
    batch_size=32,
    class_mode='categorical',
    subset='training'
)

# 验证数据生成器
val_generator = train_datagen.flow_from_directory(
    val_dir,
    target_size=(224, 224),
    batch_size=32,
    class_mode='categorical',
    subset='validation'
)

我们使用 MobileNetV2 作为基础模型，它在移动设备上表现良好：

from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D
from tensorflow.keras.models import Model

# 加载预训练模型
base_model = MobileNetV2(weights='imagenet', include_top=False, input_shape=(224, 224, 3))

# 添加自定义层
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(1024, activation='relu')(x)
predictions = Dense(241, activation='softmax')(x)  # 241 classes for cnfood-241

# 构建完整模型
model = Model(inputs=base_model.input, outputs=predictions)

# 冻结基础模型层
for layer in base_model.layers:
    layer.trainable = False

# 编译模型
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# 训练模型
history = model.fit(
    train_generator,
    steps_per_epoch=train_generator.samples // 32,
    epochs=10,
    validation_data=val_generator,
    validation_steps=val_generator.samples // 32
)

训练完成后，我们将模型转换为 int8 量化的.tflite 格式：

import tensorflow as tf

# 转换模型为 TFLite 格式
converter = tf.lite.TFLiteConverter.from_keras_model(model)

# 设置量化参数
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = representative_data_gen

# 定义代表性数据集生成函数
def representative_data_gen():
    for input_value, _ in train_generator.take(100):
        yield [input_value]

# 执行量化转换
quantized_tflite_model = converter.convert()

# 保存量化模型
with open('quantized_model.tflite', 'wb') as f:
    f.write(quantized_tflite_model)

我们对量化前后的模型进行了对比测试，结果如下：