TensorFlow系列：第四讲：MobileNetV2实战

一. 加载数据集

编写工具类，实现数据集的加载

import keras

"""
加载数据集工具类
"""


class DatasetLoader:
    def __init__(self, path_url, image_size=(224, 224), batch_size=32, class_mode='categorical'):
        self.path_url = path_url
        self.image_size = image_size
        self.batch_size = batch_size
        self.class_mode = class_mode

    # 不使用图像增强
    def load_data(self):
        # 加载训练数据集
        train_data = keras.preprocessing.image_dataset_from_directory(
            self.path_url + '/train',  # 训练数据集的目录路径
            image_size=self.image_size,  # 调整图像大小
            batch_size=self.batch_size,  # 每批次的样本数量
            label_mode=self.class_mode,  # 类别模式：返回one-hot编码的标签
        )

        # 加载验证数据集
        val_data = keras.preprocessing.image_dataset_from_directory(
            self.path_url + '/validation',  # 验证数据集的目录路径
            image_size=self.image_size,  # 调整图像大小
            batch_size=self.batch_size,  # 每批次的样本数量
            label_mode=self.class_mode  # 类别模式：返回one-hot编码的标签
        )
        # 加载测试数据集
        test_data = keras.preprocessing.image_dataset_from_directory(
            self.path_url + '/test',  # 验证数据集的目录路径
            image_size=self.image_size,  # 调整图像大小
            batch_size=self.batch_size,  # 每批次的样本数量
            label_mode=self.class_mode  # 类别模式：返回one-hot编码的标签
        )
        class_names = train_data.class_names
        return train_data, val_data, test_data, class_names

二. 训练模型完整代码

import keras
from keras import layers

from utils.dataset_loader import DatasetLoader

"""
使用MobileNetV2，实现图像多分类
"""

# 模型训练地址
PATH_URL = '../data/fruits'
# 训练曲线图
RESULT_URL = '../results/fruits'
# 模型保存地址
SAVED_MODEL_DIR = '../saved_model/fruits'

#  图片大小
IMG_SIZE = (224, 224)
# 定义图像的输入形状
IMG_SHAPE = IMG_SIZE + (3,)
# 数据加载批次,训练轮数
BATCH_SIZE, EPOCH = 32, 16


# 训练模型
def train():
    # 实例化数据集加载工具类
    dataset_loader = DatasetLoader(PATH_URL, IMG_SIZE, BATCH_SIZE)
    train_ds, val_ds, test_ds, class_total = dataset_loader.load_data()

    # 构建 MobileNet 模型
    base_model = keras.applications.MobileNetV2(input_shape=IMG_SHAPE, include_top=False)
    # 将模型的主干参数进行冻结
    base_model.trainable = False
    model = keras.Sequential([
        layers.Rescaling(1. / 127.5, offset=-1, input_shape=IMG_SHAPE),
        # 设置主干模型
        base_model,
        # 对主干模型的输出进行全局平均池化
        layers.GlobalAveragePooling2D(),
        # 通过全连接层映射到最后的分类数目上
        layers.Dense(len(class_total), activation='softmax')
    ])
    # 编译模型
    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    # 模型结构
    model.summary()
    # 指明训练的轮数epoch，开始训练
    model.fit(train_ds, validation_data=val_ds, epochs=EPOCH)
    # 测试
    loss, accuracy = model.evaluate(test_ds)
    # 输出结果
    print('Mobilenet test accuracy :', accuracy, ',loss :', loss)
    # 保存模型 savedModel格式
    model.export(filepath=SAVED_MODEL_DIR)


if __name__ == '__main__':
    train()

训练模型输出如下：

模型结构：

训练进度：主要看最下边一行输出，一轮训练完成会显示训练集和验证集的正确率。

验证正确率：

保存的模型：

三. 函数式调用方式

以后的所有讲解，都基于函数式方式进行，因为函数式调用比较灵活。

# 函数式调用方式
def train1():
    # 实例化数据集加载工具类
    dataset_loader = DatasetLoader(PATH_URL, IMG_SIZE, BATCH_SIZE)
    train_ds, val_ds, test_ds, class_total = dataset_loader.load_data()

    inputs = keras.Input(shape=IMG_SHAPE)
    # 加载预训练的 MobileNetV2 模型，不包括顶层分类器，并在 Rescaling 层之后连接
    base_model = keras.applications.MobileNetV3Large(weights='imagenet', include_top=False, input_tensor=inputs)

    # 冻结 MobileNetV2 的所有层，以防止在初始阶段进行权重更新
    for layer in base_model.layers:
        layer.trainable = False
    # 在 MobileNetV2 之后添加自定义的顶层分类器
    x = layers.GlobalAveragePooling2D()(base_model.output)
    predictions = layers.Dense(len(class_total), activation='softmax')(x)
    # 构建最终模型
    model = keras.Model(inputs=base_model.input, outputs=predictions)
    # 编译模型
    model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
    # 查看模型结构
    model.summary()
    model.fit(train_ds, validation_data=val_ds, epochs=EPOCH)
    # 测试
    loss, accuracy = model.evaluate(test_ds)
    # 输出结果
    print('Mobilenet test accuracy :', accuracy, ',loss :', loss)
    # 保存模型 savedModel格式
    model.export(filepath=SAVED_MODEL_DIR)

四. 保存训练过程曲线图

在训练模型时，我们不可能时时盯着训练数据结果，如果把训练过程曲线保存成图片，这样就比较方便查看。

在项目中编写一个工具类如下：

上边代码简单改造：

    # 训练模型
    history = model.fit(train_ds, validation_data=val_ds, epochs=EPOCH)
    # 保存曲线图
    Utils.trainResult(history, RESULT_URL)

曲线图如下：训练集和验证集准确率上升，损失率下降，这是完美的表现。

五. 模型可视化批量测试

编写可视化批量测试工具类：

import keras
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.patches import FancyBboxPatch

from utils.dataset_loader import DatasetLoader

"""
模型工具类
"""


class ModelUtil:
    def __init__(self, saved_model_dir, path_url):
        self.save_model_dir = saved_model_dir  # savedModel 模型保存地址
        self.path_url = path_url  # 模型训练数据地址

    # 批量识别 进行可视化显示
    def batch_evaluation(self, class_mode='categorical', image_size=(224, 224), num_images=25):
        dataset_loader = DatasetLoader(self.path_url, image_size=image_size, class_mode=class_mode)
        train_ds, val_ds, test_ds, class_names = dataset_loader.load_data()
        # 加载savedModel模型
        tfs_layer = keras.layers.TFSMLayer(self.save_model_dir)
        # 创建一个新的 Keras 模型，包含 TFSMLayer
        model = keras.Sequential([
            keras.Input(shape=image_size + (3,)),  # 根据你的模型的输入形状
            tfs_layer
        ])

        plt.figure(figsize=(10, 10))
        for images, labels in test_ds.take(1):
            # 使用模型进行预测
            outputs = model.predict(images)
            for i in range(num_images):
                plt.subplot(5, 5, i + 1)
                image = np.array(images[i]).astype("uint8")
                plt.imshow(image)
                index = int(np.argmax(outputs[i]))
                prediction = outputs[i][index]
                percentage_str = "{:.2f}%".format(prediction * 100)
                plt.title(f"{class_names[index]}: {percentage_str}")
                plt.axis("off")
        plt.subplots_adjust(hspace=0.5, wspace=0.5)
        plt.show()

使用工具类：

if __name__ == '__main__':
    # train()
    model_util = ModelUtil(SAVED_MODEL_DIR, PATH_URL)
    model_util.batch_evaluation()