[深度學習]Keras利用VGG進行遷移學習模板

# -*- coding: UTF-8 -*-
import keras
from keras import Model
from keras.applications import VGG16
from keras.callbacks import TensorBoard, ModelCheckpoint
from keras.layers import Flatten, Dense, Dropout, GlobalAveragePooling2D
from keras.models import load_model
from keras.preprocessing import image
from PIL import ImageFile
import numpy as np
import tensorflow as tf
from keras.preprocessing.image import ImageDataGenerator
from datetime import datetime
TIMESTAMP = "{0:%Y-%m-%dT%H-%M-%S/}".format(datetime.now())
ImageFile.LOAD_TRUNCATED_IMAGES = True


EPOCHS = 30
BATCH_SIZE = 16
DATA_TRAIN_PATH = 'D:/data/train'



def Train():
    #-------------準備數據--------------------------
    #數據集目錄應該是 train/LabelA/1.jpg  train/LabelB/1.jpg這樣
    gen = ImageDataGenerator(rescale=1. / 255)  
    train_generator = gen.flow_from_directory(DATA_TRAIN_PATH, (224,224)), shuffle=False,
                                              batch_size=BATCH_SIZE, class_mode='categorical')

    #-------------加載VGG模型而且添加本身的層----------------------
    #這裏本身添加的層須要不斷調整超參數來提高結果，輸出類別更改softmax層便可

    #參數說明：inlucde_top:是否包含最上方的Dense層，input_shape：輸入的圖像大小(width,height,channel)                                         
    base_model = VGG16(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
    x = base_model.output
    x=Flatten()(x)
    x = Dense(256, activation='relu')(x)
    x = Dropout(0.5)(x)
    x = Dense(1, activation='sigmoid')(x)
    predictions = Dense(2, activation='softmax')(x)
    model = Model(input=base_model.input, output=predictions)


    #-----------控制須要FineTune的層數，不FineTune的就直接凍結
    for layer in base_model.layers:
        layer.trainable = False

    #----------編譯，設置優化器，損失函數，性能指標
    model.compile(optimizer='rmsprop',
                  loss='binary_crossentropy', metrics=['accuracy'])
 
    #----------設置tensorboard,用來觀察acc和loss的曲線---------------
    tbCallBack = TensorBoard(log_dir='./logs/' + TIMESTAMP,  # log 目錄
                             histogram_freq=0,  # 按照何等頻率（epoch）來計算直方圖，0爲不計算
                             batch_size=16,  # 用多大量的數據計算直方圖
                             write_graph=True,  # 是否存儲網絡結構圖
                             write_grads=True,  # 是否可視化梯度直方圖
                             write_images=True,  # 是否可視化參數
                             embeddings_freq=0,
                             embeddings_layer_names=None,
                             embeddings_metadata=None)

    #---------設置自動保存點，acc最好的時候就會自動保存一次，會覆蓋以前的存檔---------------
    checkpoint = ModelCheckpoint(filepath='HatNewModel.h5', monitor='acc', mode='auto', save_best_only='True')

    #----------開始訓練---------------------------------------------
    model.fit_generator(generator=train_generator,
                        epochs=EPOCHS,
                        callbacks=[tbCallBack，checkpoint],
                        verbose=2
                        )


#-------------預測單個圖像--------------------------------------
def Predict(imgPath):
    model = load_model(SAVE_MODEL_NAME)
    img = image.load_img(imgPath, target_size=(224, 224))
    x = image.img_to_array(img)
    x = np.expand_dims(x, axis=0)
    res = model.predict(x)
    print(np.argmax(res, axis=1)[0])

 

以上運行環境：

Keras2.1.4

Tensorflow-gpu 1.5

CUDA9.0

cudnn7.0

python3.5