【Keras】使用预训练的模型进行图像分类fine-tuning

这里我实验用的数据集有三个类别，每个类别200张图片用于train，50张图片用于validation。

使用VGG16在imagenet上的预训练模型，keras里都有提供，本地没有会自动从网上下载。

from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
from keras.preprocessing.image import ImageDataGenerator, load_img
from keras.applications import VGG16
from keras import models
from keras import layers
from keras import optimizers

train_dir = './test-dataset/train'
validation_dir = './test-dataset/validation'
image_size = 224

一、冻结所有层

这里冻结了除top层的所有层。

vgg_conv = VGG16(weights='imagenet', include_top=False, input_shape=(image_size, image_size, 3))
for layer in vgg_conv.layers[:]:
    layer.trainable = False
for layer in vgg_conv.layers:
    print(layer, layer.trainable)

创建模型，添加新层。

model = models.Sequential()
model.add(vgg_conv)
model.add(layers.Flatten())
model.add(layers.Dense(1024, activation='relu'))
model.add(layers.Dropout(0.5))
model.add(layers.Dense(3, activation='softmax'))

model.summary()

创建data generator。这里没有使用data augmentation。

train_datagen = ImageDataGenerator(rescale=1./255)
validation_datagen = ImageDataGenerator(rescale=1./255)

train_batchsize = 50
val_batchsize = 10

train_generator = train_datagen.flow_from_directory(
    train_dir,
    target_size=(image_size, image_size),
    batch_size=train_batchsize,
    class_mode='categorical'
)
validation_generator = validation_datagen.flow_from_directory(
    validation_dir,
    target_size=(image_size, image_size),
    batch_size=val_batchsize,
    class_mode='categorical',
    shuffle=False
)

编译并训练。

model.compile(loss='categorical_crossentropy',
              optimizer=optimizers.RMSprop(lr=1e-4),
              metrics=['acc'])
history = model.fit_generator(
    train_generator,
    steps_per_epoch=train_generator.samples/train_generator.batch_size,
    epochs=20,
    validation_data=validation_generator,
    validation_steps=validation_generator.samples/validation_generator.batch_size,
    verbose=1
)
model.save('all_freezed.h5')

画出学习曲线。

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='Traning acc')
plt.plot(epochs, val_acc, 'r', label='Validation acc')
plt.title('Training and validation accuracy')
plt.legend()

plt.figure()
plt.plot(epochs, loss, 'b', label='Traning loss')
plt.plot(epochs, val_loss, 'r', label='Validation loss')
plt.title('Training and validation loss')
plt.legend()
plt.show()

二、冻结除最后三个卷积层的所有层

就是将freeze的部分改成：

vgg_conv = VGG16(weights='imagenet', include_top=False, input_shape=(image_size, image_size, 3))
for layer in vgg_conv.layers[:-4]:
    layer.trainable = False

for layer in vgg_conv.layers:
    print(layer, layer.trainable)

现在的学习曲线。

三、冻结除最后三个卷积层的所有层并且做data augmentation

在（二）的基础上，修改ImageDataGenerator，

train_datagen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=20,
    width_shift_range=0.2,
    height_shift_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest')

并且由于做了data augmentation，数据量变多，因此修改steps_per_epoch和epochs数量。

history = model.fit_generator(
    train_generator,
    steps_per_epoch=2*train_generator.samples/train_generator.batch_size,
    epochs=40,
    validation_data=validation_generator,
    validation_steps=validation_generator.samples/validation_generator.batch_size,
    verbose=1
)

学习曲线。