使用MobileNetV2训练自己的猫狗识别模型

745 阅读2分钟

本文已参与「新人创作礼」活动,一起开启掘金创作之路。

首先,需要了解一下什么是MobileNetV2,因为本博客与这个MobileNetV2有关,因为网络的训练使用的就是这个东西,它分为V1和v2两个版本,对于这个网络的介绍我就不详细讲述了,网上的说明很多,这篇博文主要讲述实战,在冻结内部层的情况下训练并保存我们自己的模型,代码我已经测试过了,使用迁移学习的方式训练准确率达到了98%多,已经符合我们日常需求了,可以将保存的模型嵌入到QT中制作一款简单的桌面级小动物识别程序,非常Nice,话不多说 下面直接上代码,在imageNet上训练识别猫和狗的代码如下

import tensorflow as tf
import numpy as np
import tensorflow_datasets as tfds
from tensorflow import keras, optimizers
import matplotlib.pyplot as plt

# 加载数据集
(raw_train, raw_validation, raw_test), metadata = tfds.load(
    'cats_vs_dogs',
    split=['train[:80%]', 'train[80%:90%]', 'train[90%:]'],
    with_info=True,
    as_supervised=True,
)


# 查看数据集的信息
get_label_info = metadata.features['label'].int2str

# 查看数据集的图片
for image, label in raw_train.take(10):
    plt.figure()
    plt.imshow(image)
    plt.title(get_label_info(label))
# 显示十张图片
# plt.show()


# 处理数据集,大小调整为(160x160),颜色调整为[-1,1]之间
def Process_dataset(image, label):
    # 转换数据格式
    image = tf.cast(image, tf.float32)
    image = (image / 127.5) - 1
    image = tf.image.resize(image, (160, 160))
    return image, label


# 使用map方法将此函数应用于数据集中的每一个项
train = raw_train.map(Process_dataset)
validation = raw_validation.map(Process_dataset)
test = raw_test.map(Process_dataset)

# 对训练集进行随机打乱和批处理操作,对验证机和测试集进行批处理操作
train_batch = train.shuffle(1000).batch(32)
validation_batch = validation.batch(32)
test_batch = test.batch(32)

# 查看一个batch_size结果,检查输出维度是否正确
for image_batch, label_batch in train_batch.take(1):
    pass
print(image_batch)

# MobileNet V2轻量级网络, 此处我们使用了自己的顶层,省略原顶层;冻结所有的内部层
in_shape = (160, 160, 3)
base_model = keras.applications.MobileNetV2(input_shape=in_shape, include_top=False, weights='imagenet')
base_model.trainable = False
base_model.summary()

# ModileNetV2,将每个(160, 160, 3)通道的输入转化为(5, 5, 1280)的特征快,
feature_batch = base_model(image_batch)
print(feature_batch.shape)

# 使用GlobalAveragePooling2D函数对5x5空间位置进行平均
global_average_layer = keras.layers.GlobalAveragePooling2D()
features_catch_average = global_average_layer(feature_batch)
print(features_catch_average.shape)

predict_layer = keras.layers.Dense(1)
predict_batch = predict_layer(features_catch_average)
print(predict_batch.shape)

model = keras.Sequential(
    [
        base_model,
        global_average_layer,
        predict_layer
    ]
)

# 使用RMSProp优化器编译模型
base_learning_rate = 0.0001
model.compile(optimizer=optimizers.RMSprop(lr=base_learning_rate), loss="binary_crossentropy", metrics=['accuracy'])

split_weight = (8, 1, 1)
num_train, num_val, num_test = (
    metadata.splits['train'].num_examples*weight/10 for weight in split_weight
)
initial_epochs = 10
steps_per_epoch = round(num_train)//64
validation_steps = 4
loss, accuracy = model.evaluate(validation_batch, steps=validation_steps)
history = model.fit(train_batch, epochs=initial_epochs, validation_data=validation_batch)
model.save("detectdogorcat.h5")

这几天忙着准备复试,博客写的稍微有点仓促,等复试完后面的博客将仔细的撰写,一起学习一起进步。