2. 自定义图像分类器

任务目标

训练一个可识别猫狗图像的简单CNN分类器，将其转换为TensorFlow Lite格式，并提供一个可在移动端或PC端运行的推理脚本，验证模型能对输入图像输出正确的类别概率。

环境依赖

  • TensorFlow 2.8+（含TFLite Converter）
  • Python 3.8+
  • NumPy、Pillow（用于图像预处理）
  • 数据集：使用TensorFlow Datasets中的cats_vs_dogs（或自行准备猫狗图片各200张以上）

完整实现代码

import tensorflow as tf
import numpy as np
from PIL import Image
import pathlib

# ========== 1. 数据准备 ==========
# 使用TFDS加载猫狗数据集（首次运行会自动下载约1GB数据）
import tensorflow_datasets as tfds
(ds_train, ds_test), info = tfds.load(
    'cats_vs_dogs',
    split=['train[:80%]', 'train[80%:]'],
    with_info=True,
    as_supervised=True
)

IMG_SIZE = 128
BATCH_SIZE = 32

def preprocess(image, label):
    image = tf.image.resize(image, (IMG_SIZE, IMG_SIZE))
    image = tf.cast(image, tf.float32) / 255.0
    return image, label

ds_train = ds_train.map(preprocess).batch(BATCH_SIZE).prefetch(1)
ds_test = ds_test.map(preprocess).batch(BATCH_SIZE).prefetch(1)

# ========== 2. 构建简单CNN模型 ==========
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(16, (3,3), activation='relu', input_shape=(IMG_SIZE, IMG_SIZE, 3)),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Conv2D(32, (3,3), activation='relu'),
    tf.keras.layers.MaxPooling2D(2, 2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])

model.compile(optimizer='adam',
              loss='binary_crossentropy',
              metrics=['accuracy'])

# ========== 3. 训练（仅演示用，实际训练需更多epoch） ==========
model.fit(ds_train, epochs=3, validation_data=ds_test)

# ========== 4. 转换为TFLite ==========
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

# 保存模型文件
tflite_path = 'cats_dogs_classifier.tflite'
with open(tflite_path, 'wb') as f:
    f.write(tflite_model)
print(f'TFLite模型已保存至: {tflite_path}')

# ========== 5. 推理验证 ==========
def load_and_preprocess_image(image_path):
    img = Image.open(image_path).resize((IMG_SIZE, IMG_SIZE))
    img_array = np.array(img, dtype=np.float32) / 255.0
    return np.expand_dims(img_array, axis=0)

# 加载TFLite解释器
interpreter = tf.lite.Interpreter(model_path=tflite_path)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 测试一张图片（请替换为实际图片路径）
test_image_path = 'path/to/your/test_image.jpg'  # 修改为实际路径
input_data = load_and_preprocess_image(test_image_path)

interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output = interpreter.get_tensor(output_details[0]['index'])

# 输出结果：sigmoid输出>0.5为狗，<0.5为猫
probability_dog = float(output[0][0])
print(f'预测为狗的概率: {probability_dog:.4f}')
print(f'预测结果: {"狗" if probability_dog > 0.5 else "猫"}')

运行与验证方法

  1. 训练模型：直接运行上述Python脚本。首次运行会下载cats_vs_dogs数据集（约1GB），训练3个epoch约需5-10分钟（视CPU/GPU性能而定）。训练完成后会在当前目录生成cats_dogs_classifier.tflite文件。
  2. 验证推理：将脚本末尾的test_image_path改为本地一张猫或狗图片的路径，重新运行脚本。观察输出概率：若图片为狗，输出应接近1.0；若为猫，应接近0.0。
  3. 通过条件：模型文件成功生成且推理脚本能对测试图片输出0~1之间的概率值，且对明显猫/狗图片给出正确分类方向（猫<0.5<狗），即视为部署链路验证通过。

移动端接入说明

生成的.tflite文件可直接用于Android（通过TensorFlow Lite Java API）或iOS（通过Swift API）。输入需为128×128×3的RGB图像，像素值归一化到[0,1]，输出为单个sigmoid值。移动端无需包含训练代码，仅需加载此模型文件并执行推理。

常见问题

  • 数据集下载失败：可手动下载猫狗图片并改用tf.keras.preprocessing.image_dataset_from_directory加载。
  • 训练准确率低：3个epoch仅用于演示流程，实际使用建议训练10个epoch以上，或使用数据增强。
  • TFLite推理结果与Keras不一致：检查预处理是否一致（尺寸、归一化方式），TFLite默认输入为NCHW格式时需调整。
← 上一章 📖 返回目录 下一章 →