keras+tensorflow实现迁移学习

coding=utf-8

“”"
author:lei
function:迁移学习
“”"

import tensorflow as tf
from tensorflow.python import keras
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
from tensorflow.python.keras.applications.vgg16 import VGG16, preprocess_input
import numpy as np

class TransforModel(object):

def __init__(self):
    # 定义训练和测试图片的变化方法，标准化和数据增强
    self.train_generator = ImageDataGenerator(rescale=1.0/255.0)  # 标准化，其他自定义
    self.test_gengerator = ImageDataGenerator(rescale=1.0/255.0)

    # 指定训练数据和测试数据的目录
    self.train_dir = "./data/train"
    self.test_dir = "./data/test"

    # 定义图片训练相关训练参数
    self.image_size = (224, 224)
    self.batch_size = 32

    # 定义迁移学习的基类模型  include_top=不包含最后的全连接层的参数
    self.base_model = VGG16(include_top=False)

    self.label_dict = {'0': 'bus', '1': 'dinosaurs', '2': 'elephants', '3': 'flowers', '4': 'horse'}

def get_local_data(self):
    """
        读取本地的图片数据以及类别
    :return: 返回训练数据和测试数据迭代器
    """
    # 使用flow_from_derectory
    train_gen = self.train_generator.flow_from_directory(self.train_dir, target_size=self.image_size, batch_size=self.batch_size, class_mode="binary", shuffle=True)

    test_gen = self.test_gengerator.flow_from_directory(self.test_dir, target_size=self.image_size, batch_size=self.batch_size, class_mode="binary", shuffle=True)

    return train_gen, test_gen

def refine_base_model(self):
    """
        微调vgg结构，5个blocks后面+全局平均池化（减少迁移学习的参数数量）+ 两个全连接层
    :return:
    """
    # 1、获取原notop模型的输出
    x = self.base_model.outputs[0]
    # print(x)  shape=(?, ?, ?, 512)

    # 2、再输出后面增加我们自己定义的结构  先进行全局池化
    # [?, ?, ?, 512] --> [?, 1*1*512]
    x = keras.layers.GlobalAveragePooling2D()(x)

    # 3、先定义新的迁移模型  上一层的输入要输入当中
    x = keras.layers.Dense(1024, activation=tf.nn.relu)(x)
    y_predict = keras.layers.Dense(5, activation=tf.nn.softmax)(x)

    # model定义新模型
    transfer_model = keras.models.Model(inputs=self.base_model.inputs, outputs=y_predict)

    return transfer_model

def freeze_model(self):
    """
        冻结vgg模型
        冻结vgg的多少，根据你的数据量
    :return:
    """
    # self.base_model.layers 获取所有层，返回层的列表
    for layer in self.base_model.layers:
        layer.trainable = False

def compile(self, model):
    """
        编译模型
    :return:
    """
    model.compile(optimizer=keras.optimizers.Adam(), loss=keras.losses.sparse_categorical_crossentropy, metrics=["accuracy"])

    return None

def fit(self, model, train_gen, test_gen):
    """
        训练模型  model.fit_generator() 而不是model.fit()
    :return:
    """
    # 保存文件
    model_ckpt = keras.callbacks.ModelCheckpoint("./ckpt/transfer_{epoch:02d}--{val_acc:.2f}.h5", monitor="val_acc", save_weights_only=True, save_best_only=True, mode="auto", period=1)

    model.fit(train_gen, epochs=3, validation_data=test_gen, callbacks=[model_ckpt])

    return None

def predict(self, model):
    """
        预测类别
    :return:
    """
    # 加载模型，transfer_model
    model.load_weights("./ckpt/transfer_03--0.95.h5")

    # 读取图片，处理
    image = load_img("./bus/300.jpg", target_size=(224, 224))
    image = img_to_array(image)
    print(image.shape)

    # 转换成4维函数 (224, 224, 3) --> (1, 224, 224, 1)
    image = image.reshape([1, image.shape[0], image.shape[1], image.shape[2]])
    print(image)
    # model.predict()

    # 预测结果进行处理
    image = preprocess_input(image)
    predictions = model.predict(image)
    print(predictions)

    index = np.argmax(predictions, axis=1)

    # 返回预测的结果
    print(self.label_dict[str(index[0])])

if name == ‘main’:
tm = TransforModel()
# 训练
# train_gen, test_gen = tm.get_local_data()
# print(train_gen)
# print(test_gen)
# # 通过迭代器 打印
# # for data in train_gen:
# # print(data[0].shape, data[1].shape)
# print(tm.base_model.summary())
#
# model = tm.refine_base_model()
# print(model)
# # 冻结模型
# tm.freeze_model()
#
# tm.compile(model)
#
# tm.fit(model, train_gen, test_gen)

# 测试
model = tm.refine_base_model()
tm.predict(model)