ESP-DL部署魔改MobilenetV1—1. 模型训练

作为入门，以分类问题为例，对mobilenetV1进行修改，设计了一个简单的深度学习模型，实现对猫狗的分类。在此之前，需要您首先具备一些基础的深度学习相关知识。

软件准备

• tensorflow2
• tf2onnx

数据集处理

数据集准备

针对这个分类问题，我们直接使用开源的kaggle猫狗识别数据集进行训练。该数据集分为猫狗两类，首先我们需要对数据集进行整理，从而方便tensorflow读取。
我们将文件格式修改为如下：

cat-vs-dog/
|-- cat/
|   |-- cat_001.jpg
|   |-- cat_002.jpg
|   |-- cat_003.jpg
|   |-- ...
|-- dog/
    |-- dog_001.jpg
    |-- dog_002.jpg
    |-- dog_003.jpg
    |-- ...

这样我们就完成了数据集的准备。
数据集读取
我们使用image_dataset_from_directory来读取数据集，首先导入该包，并定义数据集的路径和batch size的大小。

from tensorflow.keras.preprocessing import image_dataset_from_directory

定义数据集的路径

data_dir = '/tmp/code/dataset/cat_vs_dog'
batch_size = 128

之后就能够直接导入我们准备的数据集并将其划分为训练集和测试集了。其中，训练数据集用于训练模型、测试数据集用于测试模型的泛化性能。

加载数据集

train_dataset = image_dataset_from_directory(
    data_dir,
    labels='inferred',
    label_mode='categorical',
    class_names=None,
    color_mode='rgb',
    batch_size=batch_size,
    image_size=(96, 96),
    shuffle=True,
    seed=123,
    validation_split=0.2,  # 20% 的数据用于验证
    subset='training',  # 加载训练集
    interpolation='bilinear',
    follow_links=False
)

加载测试集

validation_dataset = image_dataset_from_directory(
    data_dir,
    labels='inferred',
    label_mode='categorical',
    class_names=None,
    color_mode='rgb',
    batch_size=batch_size,
    image_size=(96, 96),
    shuffle=True,
    seed=123,
    validation_split=0.2,  # 20% 的数据用于测试
    subset='validation',  # 加载测试集
    interpolation='bilinear',
    follow_links=False
)

加载完成后会显示数据集的相关信息。
其中我们将训练集和测试集按照8：2的比例进行划分，并且将图片统一resize为96*96大小，rgb格式。其中还有一个重要的参数为label_mode，它指定返回的标签格式。 如果是二分类只有一个输出，可以设为‘binary’；如果输出为one-hot 编码的标签需要设为 ‘categorical’。两个的具体区别建议百度，这里不具体阐述了。

读取成功后，就会分别输出我们训练集和测试集的数据大小。

数据增强

我们读取完数据集图片后，我们需要对图片进行处理，如归一化等操作，因此我们可以定义一个处理函数，用于对数据集进行处理

定义数据增强函数

我们定义如下函数，将图片数据归一化为[0,1]之间的数据。

def data_augmentation(image, label):
    image = image/255
    # image = tf.image.random_flip_left_right(image)
    # image = tf.image.random_brightness(image, max_delta=0.1)
    # image = tf.image.random_contrast(image, lower=0.9, upper=1.1)
    return image, label

将数据增强应用于训练数据集

之后我们将这个函数运用于我们前面读取的数据集

train_dataset = train_dataset.map(data_augmentation)
validation_dataset = validation_dataset.map(data_augmentation)

数据查看

在数据集准备和读取完成后，我们来查看以下数据以及维度大小是否符合预期，并输出一张数据集中的图片看一下是否正常。

# 查看数据集的元素
for images, labels in train_dataset.take(1):
    cal_images = images[0]
    print(images.shape)  # 输出图像的形状
    print(labels.numpy())  # 输出标签
    print(images[0])  # 输出图像
 
import numpy as np
from PIL import Image

data = cal_images  # 示例数据
image_data = (data * 255).astype(np.uint8)# 将值范围从 (0, 1) 调整到 (0, 255)
img = Image.fromarray(image_data, 'RGB')# 创建 PIL 图像对象
img.show()# 显示图像

模型搭建

在完成数据集的准备后，我们需要进行模型的搭建，这里我们以keras中的mobilenet为例，进行模型的搭建

import tensorflow as tf
from mobilenet import MobileNet

base_model = MobileNet(
    input_shape=(96, 96, 3),
    alpha=1.0,
    depth_multiplier=1,
    dropout=0.1,
    include_top=False,
    # weights='imagenet',
    input_tensor=None,
    pooling=None,
    classes=2,
    classifier_activation='softmax'
)

base_model.summary()

虽然mobilenet是轻量级模型，但是其参数量对于esp32来说还是有点大，而模型太大不仅会导致模型无法部署，也会导致推理速度很慢，因此我们基于keras中的mobilenet，对其进行修改精简，得到新的更小的模型并保存在mobilenet.py文件中，便于直接引用。

此外，由于我们这里导入的模型并不包含最后分类的层，下面我们自行将最后几层添加进来

classes = 2

# 添加新的层用于 2 类别的分类任务
x = base_model.output
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dropout(0.5)(x)  # 添加 Dropout 层
x = tf.keras.layers.Dense(2)(x)# 使用 Dense 层但不带 softmax 激活
predictions = tf.keras.layers.Softmax()(x)# 添加 Softmax 层

# 创建新的模型
model = tf.keras.Model(inputs=base_model.input, outputs=predictions)

# 查看模型结构
model.summary()

其中的Dropout层是为了防止模型过拟合而添加的，一般位于全连接层前面，其参数可以根据需要自行设置。

模型训练

在tensorflow中，训练模型非常简单，我们首先先对模型参数进行设置

# 定义指数衰减的学习率
steps_per_epoch = 20000//batch_size

initial_learning_rate = 0.01  # 初始学习率
decay_steps = steps_per_epoch*20  # 每 decay_steps 步后学习率衰减
decay_rate = 0.95  # 学习率衰减的比例
staircase = True  # 是否使学习率呈阶梯状下降

lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=initial_learning_rate,
    decay_steps=decay_steps,
    decay_rate=decay_rate,
    staircase=staircase
)

optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)

这里我们根据训练集的大小设置了每个epoch需要训练几步，并且设置了初始学习率，并令学习率在训练过程中呈指数衰减。
设置完成后，运行如下命令就能够启动模型的训练了

model.compile(optimizer = optimizer, loss='categorical_crossentropy', metrics=['acc'])

# 训练模型
history = model.fit(
    train_dataset,
    epochs=100,
    validation_data=validation_dataset,
    verbose=2
)

模型导出

在模型训练完成后，我们运行如下命令对将训练好的模型进行保存

# 保存模型为 SavedModel 格式
save_path = 'saved_model'
tf.saved_model.save(model, save_path)

此外为了导出tf模型为onnx格式，我们还需要安装tf2onnx库

pip install -U tf2onnx

之后导出tf模型为onnx格式，分别有命令行和python代码两种方式

python -m tf2onnx.convert --saved-model './saved_model' --output model.onnx --opset 13

import tf2onnx
import tensorflow as tf

# 将模型转换为ONNX
input_dim = (240, 240, 3)  # 输入图像的尺寸
spec = (tf.TensorSpec((None, *input_dim), tf.float32, name="input"),)
output_path = "model.onnx"

# 转换模型
with tf.device('/CPU:0'):  # 使用CPU进行转换
    # 调用 convert 方法进行转换
    model_proto, _ = tf2onnx.convert.from_keras(
        model,# 需要保存的模型
        input_signature=spec,
        opset=13,
        output_path=output_path
    )