深度探究TensorFlow2.0版本图像分类进阶指南

2023-11-21 06:18:00

前言

随着深度学习的飞速发展，图像分类任务在计算机视觉领域中变得越来越重要。本文将详细介绍如何使用TensorFlow2.0版本来构建一个图像分类模型。我们将使用猫狗大战数据集来训练模型，并使用ReduceLROnPlateau来调整学习率。最后，我们将打印loss结果生成jpg图片，以直观地展示模型的训练过程。

数据集准备

本文中，我们将使用猫狗大战数据集来训练图像分类模型。该数据集包含1万张图像，其中包含5000张猫的图像和5000张狗的图像。我们将使用ImageDataGenerator类来加载和预处理图像数据。

import tensorflow as tf

# 加载猫狗大战数据集
data_dir = 'path/to/cat_dog_dataset'
train_ds = tf.keras.preprocessing.image_dataset_from_directory(
    data_dir,
    batch_size=32,
    image_size=(224, 224),
    shuffle=True,
)

# 预处理图像数据
train_ds = train_ds.map(lambda x, y: (tf.image.resize(x, (224, 224)), y))
train_ds = train_ds.map(lambda x, y: (tf.image.random_flip_left_right(x), y))
train_ds = train_ds.map(lambda x, y: (tf.image.random_rotation(x, 0.2), y))
train_ds = train_ds.map(lambda x, y: (tf.image.random_zoom(x, (0.8, 1.2)), y))

模型构建

我们将使用自定义的CNN网络来构建图像分类模型。该网络包含一个卷积层、一个池化层、一个全连接层和一个输出层。

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(224, 224, 3)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(2, activation='softmax')
])

模型编译

我们将使用二分类交叉熵损失函数和Adam优化器来编译模型。

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

模型训练

我们将使用ReduceLROnPlateau回调函数来调整学习率。该回调函数将在验证集准确率不再提高时降低学习率。

callbacks = [
    tf.keras.callbacks.ReduceLROnPlateau(patience=3)
]

model.fit(train_ds, epochs=10, callbacks=callbacks)

模型评估

我们将使用测试集来评估模型的性能。

test_dir = 'path/to/test_dataset'
test_ds = tf.keras.preprocessing.image_dataset_from_directory(
    test_dir,
    batch_size=32,
    image_size=(224, 224),
    shuffle=True,
)

test_ds = test_ds.map(lambda x, y: (tf.image.resize(x, (224, 224)), y))

loss, accuracy = model.evaluate(test_ds)

print('Test loss:', loss)
print('Test accuracy:', accuracy)

可视化结果

我们将使用matplotlib库来打印loss结果生成jpg图片。

import matplotlib.pyplot as plt

plt.plot(model.history['loss'])
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Loss curve')
plt.savefig('loss.jpg')