TensorFlow Keras ImageDataGenerator VGG16 Finetune

#%%

from keras.applications.vgg16 import VGG16
from keras.preprocessing import image
from keras.preprocessing.image import ImageDataGenerator
import numpy as np
from keras.models import Sequential
from keras.layers import Dropout, Flatten, Dense
from keras.optimizers import SGD

#%%

vgg16_model = VGG16(weights='imagenet', include_top=False, input_shape=(150, 150, 3))

#%%

# 全连接层
top_model = Sequential()
top_model.add(Flatten(input_shape=vgg16_model.output_shape[1:]))
top_model.add(Dense(256, activation='relu'))
top_model.add(Dropout(0.5))
top_model.add(Dense(2, activation='softmax'))

#%%

# 载入训练过的权值
top_model.load_weights('bottleneck_fc_model.h5')

#%%

# VGG16 和全连接层连起来
model = Sequential()
model.add(vgg16_model)
model.add(top_model)

#%%

# 训练集数据生成器
train_datagen = ImageDataGenerator(
    rescale=1. / 255,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True
)
# 测试集数据生成器
test_dategen = ImageDataGenerator(
    rescale=1. / 255
)

#%%

batch_size = 32
# 训练集生成器
train_generator = train_datagen.flow_from_directory(
    # 训练集路径
    'image/train',
    # 设置生成的图片大小
    target_size=(150, 150),
    # 批次大小
    batch_size=batch_size
)
# 测试集生成器
test_generator = test_dategen.flow_from_directory(
    # 训练集路径
    'image/test',
    # 设置生成的图片大小
    target_size=(150, 150),
    # 批次大小
    batch_size=batch_size
)

#%%

model.compile(loss='categorical_crossentropy',
              # 梯度下降
              optimizer=SGD(lr=1e-4, momentum=0.9),
              metrics=['accuracy'])

#%%

model.fit_generator(
    train_generator,
    steps_per_epoch=train_generator.samples / batch_size,
    epochs=2,
    validation_data=test_generator,
    validation_steps=test_generator.samples / batch_size
)

#%%

# 模型保存
model.save('vgg16_finetune_cat_dog.h5')