一. 加载数据集
编写工具类，实现数据集的加载
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()