基于VGG16的迁移学习-训练自己的数据集

简介
本程序基于VGG16的进行迁移学习，通过训练自己的数据集实现三分类，通过冻结原网络的特征提取层，使卷积层和池化层的权重保持不变。由于待分类的数据与原先 VGG16 的分类数据不同，删除原来的全连接层，在特征提取层之后添加全局平均池化层(Global Average pooling），再增加两个全新的全连接层（Fully-connected layer ），最后一层全连接层分类数与数据集的类数相匹配，通过重新训练确定最后几层的参数信息，来实现分类目标。并用plot画出Training acc,validation acc,Training loss,validation loss。基于VGG16的迁移学习模型如图所示。![在这里插入图片描述](https://img-blog.csdnimg.cn/20210325184301245.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L1
Nvbmd5b25nY2hhbzE5OTU=,size_16,color_FFFFFF,t_70)

import os
import numpy as np
from keras import Model
from keras.applications import VGG16
import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing import image
from keras.layers import GlobalMaxPooling2D, Flatten, Dense, Dropout, GlobalAveragePooling2D
import random

def DataSet():
    train_path_A = 'C:/Users/HASee/Desktop/shilu/A/'
    train_path_B = 'C:/Users/HASee/Desktop/shilu/B/'
    train_path_C = 'C:/Users/HASee/Desktop/shilu/C/'

    imglist_train_A = os.listdir(train_path_A)
    imglist_train_B = os.listdir(train_path_B)
    imglist_train_C = os.listdir(train_path_C)

    X_train = np.empty((len(imglist_train_A) + len(imglist_train_B)+ len(imglist_train_C), 224, 224, 3))
    Y_train = np.empty((len(imglist_train_A) + len(imglist_train_B)+ len(imglist_train_C), 3))
    count = 0
    for img_name in imglist_train_A:
        img_path = train_path_A + img_name
        img = image.load_img(img_path, target_size=(224, 224))
        img = image.img_to_array(img) / 255.0

        X_train[count] = img
        Y_train[count] = np.array((1, 0, 0))
        count += 1

    for img_name in imglist_train_B:
        img_path = train_path_B + img_name
        img = image.load_img(img_path, target_size=(224, 224))
        img = image.img_to_array(img) / 255.0

        X_train[count] = img
        Y_train[count] = np.array((0, 1, 0))
        count += 1

    for img_name in imglist_train_C:
        img_path = train_path_C + img_name
        img = image.load_img(img_path, target_size=(224, 224))
        img = image.img_to_array(img) / 255.0

        X_train[count] = img
        Y_train[count] = np.array((0, 0, 1))
        count += 1

    index = [i for i in range(len(X_train))]
    random.shuffle(index)
    X_train = X_train[index]
    Y_train = Y_train[index]

    return X_train, Y_train


X_train, Y_train = DataSet()
print('X_train shape : ', X_train.shape)
print('Y_train shape : ', Y_train.shape)


# # model
base_model = VGG16(weights='imagenet', include_top=False)   # create the base pre-trained model
x = base_model.output
x = GlobalAveragePooling2D()(x)                                   # add a global spatial average pooling layer
x = Dense(512, activation='relu')(x)                # let's add a fully-connected layer
x = Dropout(0.5)(x)
predictions = Dense(3, activation='softmax')(x)       # and a logistic layer -- let's say we have 2 classes
model = Model(inputs=base_model.input, outputs=predictions)  # this is the model we will train
for layer in base_model.layers:              # first: train only the top layers (which were randomly initialized)
    layer.trainable = False                      # i.e. freeze all convolutional InceptionV3 layers

# compile the model (should be done *after* setting layers to non-trainable)
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])


#Tensorboard = TensorBoard(log_dir="logs", histogram_freq=1)
history = model.fit(X_train, Y_train, epochs=2, batch_size=20, verbose=1, validation_split=0.33)

acc = history.history['acc']
val_acc = history.history['val_acc']
loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(len(acc))

plt.plot(epochs, acc, 'b', label='Training acc')
plt.plot(epochs, val_acc, 'r', label='validation acc')
plt.title('Training and validation acc')
plt.legend(loc='lower right')
plt.figure()

plt.plot(epochs, loss, 'b', label='Training loss')
plt.plot(epochs, val_loss, 'r', label='validation loss')
plt.title('Training and validation loss')
plt.legend()
plt.show()
# 保存模型vgg
# model.save('my_model-vgg19.h5')