基于 VGG16 的图像分类模型训练与评估

1. 简介

本文介绍了如何使用 VGG16 预训练模型进行图像分类。VGG16 是一个经典的卷积神经网络模型，在 ImageNet 数据集上取得了很好的效果。我们可以利用它提取图像的特征，并构建一个简单的分类模型进行训练和评估。

2. 环境准备

Python 3.x
Keras
NumPy
Matplotlib

3. 代码实现

# -*- coding: utf-8 -*-
import os
import matplotlib.pyplot as plt
import numpy as np
from keras.applications.vgg16 import VGG16
from keras.preprocessing.image import ImageDataGenerator

data_dir = 'ss' + os.sep + 'image'
train_dir = os.path.join(data_dir, 'train')
test_dir = os.path.join(data_dir, 'test')

vgg_model = VGG16(weights='imagenet', include_top=False, input_shape=(150, 150, 3))
vgg_model.summary()

datagen = ImageDataGenerator(rescale=1./255)

batch_size = 20

# 定义提取特征函数
def extract_features(directory, sample_count):
    # VGG16 最后池化层的输出特征初始化储存张量
    features = np.zeros(shape=(sample_count, 4, 4, 512))
    labels = np.zeros(shape=(sample_count))
    # 从数据样本文件夹按 batch_size 数加载图像
    generator = datagen.flow_from_directory(directory,
                                            target_size=(150, 150),
                                            batch_size=batch_size,
                                            class_mode='binary')
    i = 0
    for inputs_batch, labels_batch in generator:
        features_batch = vgg_model.predict(inputs_batch)
        features[i * batch_size : (i + 1) * batch_size] = features_batch
        labels[i * batch_size : (i + 1) * batch_size] = labels_batch
        i += 1
        # 
        if i * batch_size >= sample_count:
            break
    return features, labels

# 
train_features, train_labels = extract_features(train_dir, 2000)
# validation_features, validation_labels = extract_features(validation_dir, 1000)
test_features, test_labels = extract_features(test_dir, 1000)
# 
train_features = np.reshape(train_features, (2000, 4 * 4 * 512))
# validation_features = np.reshape(validation_features, (1000, 4 * 4 * 512))
test_features = np.reshape(test_features, (1000, 4 * 4 * 512))

from keras import models
from keras import layers
from keras import optimizers
# 
model = models.Sequential()
model.add(layers.Dense(256, activation='relu', input_dim=4 * 4 * 512))
model.add(layers.Dropout(0.5))
model.add(layers.Dense(1, activation='sigmoid'))
model.compile(optimizer=optimizers.Adam(lr=2e-5),
              loss='binary_crossentropy',
              metrics=['acc'])

history = model.fit(train_features, train_labels,
                    epochs=30,
                    batch_size=batch_size,
                    validation_data=(test_features,
                                     test_labels))

acc = history.history['acc']
val_acc = history.history['val_acc']
epochs = range(1, len(acc) + 1)
plt.figure()
plt.plot(epochs, acc, 'bo', label='训练集准确率')
plt.plot(epochs, val_acc, 'b', label='验证集准确率')
plt.title('训练集与验证集准确率曲线')
plt.legend()
plt.show()

# 
acc = history.history['loss']
val_acc = history.history['val_loss']
epochs = range(1, len(acc) + 1)
plt.figure()
plt.plot(epochs, acc, 'bo', label='训练集损失')
plt.plot(epochs, val_acc, 'b', label='验证集损失')
plt.title('训练集与验证集损失曲线')
plt.legend()
plt.show()

4. 总结

本文介绍了使用 VGG16 预训练模型进行图像分类的流程，包括数据预处理、特征提取、模型构建、训练和评估。并使用 Matplotlib 绘制了训练集和验证集的准确率以及损失函数曲线。

5. 注意事项

数据集的预处理和扩充是影响模型性能的重要因素。
模型的超参数设置需要根据实际情况进行调整。
可以尝试使用不同的优化器和损失函数来提高模型性能。
可以尝试使用其他预训练模型进行实验。

6. 参考链接