【深度神經網絡】四、Mini-VGG實現CIFAR10數據集分類

概要

本篇博客主要講解了CIFAR10數據集的預處理，Mini-VGG的tensorflow和keras實現，並實現了CIFAR數據集的分類。

VGG16網絡架構論文講解請詳見：【深度神經網絡】三、VGG網絡架構詳解

整個項目的github地址爲：Mini-VGG-CIFAR10 ，如果喜歡集的點個Star

---

一、 Cifar10數據集說明

爲了實現VGG16網絡對CIFAR10數據集的分類，我們首先得對CIFAR10進行一個詳細介紹

Cifar10數據集共有60000張彩色圖像，這些圖像是32*32，分爲10個類，每類6000張圖。其中，有50000張用於訓練，構成了5個訓練批，每一批10000張圖；另外10000用於測試，單獨構成一批。測試批的數據裏，取自10類中的每一類，每一類隨機取1000張。抽剩下的就隨機排列組成了訓練批。注意一個訓練批中的各類圖像並不一定數量相同，總的來看訓練批，每一類都有5000張圖。
下面這幅圖就是列舉了10各類，每一類展示了隨機的10張圖片：

該數據是由以下三個人收集而來：Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton。第一位是AlexNet的提出者，第三位就更不用說了——深度學習的奠基人。

該數據集的下載網址爲：http://www.cs.toronto.edu/~kriz/cifar.html 。這個數據主要有三個下載版本：Python、Matlab和二進制文件（適合於C語言）。由於我主要是利用tensorflow和Keras來實現VGG，因此我下載的是Python版本的數據集。從網站上可以看出，無論下載那個版本的數據集文件都不是挺大，足夠學習跑跑程序用。

下面開始導入Cifar10數據集。將官網上下載的數據集打開之後，文件結構如下圖所示。主要包含了5個data_batch文件data_batch_1至data_batch_5、1個test_batch文件和1個batches的meta文件。

從Cifar10數據集官網上的介紹來看，5個data_batch文件和test_batch文件是利用pickel序列化之後的文件因此在導入 Cifar10數據集必須利用pickel進行解壓數據，之後將數據還原。5個data_batch文件和test_batch文件分別代表5個訓練集批次和測試集，因此我們首先利用pickel編寫解壓函數：

import pickle as pk

def unpickle(data_path):
    """
    這是解壓pickle數據的函數
    :param data_path: 數據路徑
    """
    # 解壓數據
    with open(data_path,'rb') as f:
        data_dict = pk.load(f,encoding='latin1')
        # 獲取標籤，形狀(10000,)
        labels = np.array(data_dict['labels'])
        # 獲取圖像數據
        data = np.array(data_dict['data'])
        # 轉換圖像數據形狀,形狀爲(10000,32,32,3)
        data = np.reshape(data,(10000,3,32,32)).transpose(0,2,3,1).astype("float")
    return data,labels

在下面項目中訓練Mini-VGG時候使用Keras官方的ImageDataGenerator來構造數據生成器，因此首先需要將官方CIFAR10數據集存儲轉化成適應ImageDataGenerator接口的數據集格式。在這裏我們將CIFAR官方中data_batch1至data_batch5作爲訓練集，test_batch作爲驗證集，即訓練集有5萬張圖片，驗證集有1萬張圖片。

同時，6萬張圖像需要利用opencv庫重新寫入內存這會涉及大量I/O操作，因此爲了加快圖像寫入內存速度，採用了異步多進程方式來實現CIFAR10數據集的寫入內存。CIFAR10數據集格式轉化腳本如下：

# -*- coding: utf-8 -*-
# @Time    : 2020/3/24 11:39
# @Author  : Dai PuWei
# @Email   : [email protected]
# @File    : cifar_preprocess.py
# @Software: PyCharm

import os
import cv2
import numpy as np
import pickle as pk
import multiprocessing
from multiprocessing import Pool

def unpickle(data_path):
    """
    這是解壓pickle數據的函數
    :param data_path: 數據路徑
    """
    # 解壓數據
    with open(data_path,'rb') as f:
        data_dict = pk.load(f,encoding='latin1')
        # 獲取標籤，
        #labels = np.array(data_dict['fine_labels'])          # CIFAR100數據集
        labels = np.array(data_dict['labels'])                # CIFAR10數據集
        # 獲取圖像數據
        data = np.array(data_dict['data'])
        # 轉換圖像數據形狀,形狀爲(10000,32,32,3)
        size = len(data)
        data = np.reshape(data,(size,3,32,32)).transpose(0,2,3,1).astype("int32")
    return data,labels

def save_single_iamge(image,image_path):
    """
    這是保存單張圖像的函數
    :param image: 圖像
    :param image_path: 圖像地址
    :return:
    """
    print(image_path)
    cv2.imwrite(image_path,image)

def save_batch_images(batch_images,batch_image_ptahs):
    """
    這是保存批量圖像的函數
    :param batch_images: 批量圖像
    :param batch_image_ptahs: 圖量圖像地址
    :return:
    """
    for image,image_path in zip(batch_images,batch_image_ptahs):
        save_single_iamge(image,image_path)

def cifar_preprocess(cifar_dataset_dir,new_cifar_dataset_dir,batch_size):
    """
    這是對CIFAR數據集進行預處理的函數
    :param cifar_dataset_dir: cifar數據集
    :param new_cifar_dataset_dir: 新cifar數據集
    :param batch_size: 小批量數據集規模
    :return:
    """
    # 初始化路徑原始CIFAR10數據集訓練集和測試集路徑
    train_batch_paths = [os.path.join(cifar_dataset_dir,"data_batch_%d"%(i+1)) for i in range(5)]
    val_batch_path = os.path.join(cifar_dataset_dir,'test_batch')

    # 初始化新格式下CIFAR10數據集的訓練、驗證和測試集目錄
    new_train_dataset_dir = os.path.join(new_cifar_dataset_dir, "train")
    new_val_dataset_dir = os.path.join(new_cifar_dataset_dir, 'val')
    if not os.path.exists(new_cifar_dataset_dir):
        os.mkdir(new_cifar_dataset_dir)
    if not os.path.exists(new_train_dataset_dir):
        os.mkdir(new_train_dataset_dir)
    if not os.path.exists(new_val_dataset_dir):
        os.mkdir(new_val_dataset_dir)

    # 初始化訓練驗證集中，每個類別圖像的目錄
    label_names = ["airplane","automobile","bird","cat","deer","dog","frog","horse","ship","truck"]
    for label_name in label_names:
        train_label_dir = os.path.join(new_train_dataset_dir,label_name)
        val_label_dir = os.path.join(new_val_dataset_dir, label_name)
        if not os.path.exists(train_label_dir):
            os.mkdir(train_label_dir)
        if not os.path.exists(val_label_dir):
            os.mkdir(val_label_dir)

    # 解析原始數據集文件
    train_data = []
    train_labels = []
    for i,train_batch_path in enumerate(train_batch_paths):
        batch_images,batch_labels = unpickle(train_batch_path)
        train_data.append(batch_images)
        train_labels.append(batch_labels)
    train_data = np.concatenate(train_data,axis=0)
    train_labels = np.concatenate(train_labels,axis=0)
    val_data,val_labels = unpickle(val_batch_path)
    print("訓練數據集維度：",np.shape(train_data))
    print("訓練標籤集維度：",np.shape(train_labels))
    print("驗證數據集維度：",np.shape(val_data))
    print("驗證標籤集維度：",np.shape(val_labels))

    # 遍歷每種類別，分別獲取每種類別圖像，然後按照比例劃分成訓練和測試集，並生成對應圖像的路徑
    train_index = np.arange(len(train_labels))
    train_images = []
    new_train_image_paths = []
    for i,label_name in enumerate(label_names):
        # 獲取指定類別的下標
        label_index = np.random.permutation(train_index[train_labels == i])
        train_images.append(train_data[label_index])
        # 生成指定類別訓練集圖像的路徑
        batch_new_train_image_paths = []
        for j,index in enumerate(label_index):
            image_name = "%06d.jpg"%(j)
            new_train_image_path = os.path.join(new_train_dataset_dir,label_name,image_name)
            batch_new_train_image_paths.append(new_train_image_path)
        new_train_image_paths.append(np.array(batch_new_train_image_paths))
    train_images = np.concatenate(train_images,axis=0)
    new_train_image_paths = np.concatenate(new_train_image_paths,axis=0)

    # 遍歷每種類別，分別獲取每種類別的驗證圖像，並生成對應測試圖像的路徑
    val_index = np.arange(len(val_labels))
    val_images = []
    new_val_image_paths = []
    for i, label_name in enumerate(label_names):
        label_index = val_index[val_labels == i]
        val_images.append(val_data[label_index])
        batch_new_test_image_paths = []
        for j, index in enumerate(label_index):
            image_name = "%06d.jpg" % (j)
            new_val_image_path = os.path.join(new_val_dataset_dir, label_name, image_name)
            batch_new_test_image_paths.append(new_val_image_path)
        new_val_image_paths.append(np.array(batch_new_test_image_paths))
    val_images = np.concatenate(val_images,axis=0)
    new_val_image_paths = np.concatenate(new_val_image_paths,axis=0)

    print("訓練數據集維度：", np.shape(train_images))
    print("訓練標籤集維度：", np.shape(new_train_image_paths))
    print("驗證數據集維度：", np.shape(val_images))
    print("驗證標籤集維度：", np.shape(new_val_image_paths))
    print(new_train_image_paths)
    print(new_val_image_paths)

    # 按照每種類別，分別劃分成小批量訓練集，利用多進程保存圖像
    print("Start generating train dataset")
    train_size = len(new_train_image_paths)
    pool = Pool(processes=multiprocessing.cpu_count())
    for i,start in enumerate(np.arange(0,train_size,batch_size)):
        end = int(np.min([start+batch_size,train_size]))               # 防止最後一組數量不足batch_size
        batch_train_images = train_images[start:end]
        batch_new_train_image_paths = new_train_image_paths[start:end]
        #print("進程%d需要處理%d張圖像"%(i,len(batch_new_train_image_paths)))
        pool.apply_async(save_batch_images,args=(batch_train_images,batch_new_train_image_paths))
    pool.close()
    pool.join()
    print("Finish generating train dataset")

    # 按照每種類別，分別劃分成小批量訓練集，利用多進程保存圖像
    print("Start generating val dataset")
    val_size = len(new_val_image_paths)
    pool = Pool(processes=multiprocessing.cpu_count())
    for i, start in enumerate(np.arange(0, val_size, batch_size)):
        end = int(np.min([start + batch_size, train_size]))  # 防止最後一組數量不足batch_size
        batch_val_images = val_images[start:end]
        batch_new_val_image_paths = new_val_image_paths[start:end]
        #print("進程%d需要處理%d張圖像" % (i, len(batch_new_val_image_paths)))
        pool.apply_async(save_batch_images, args=(batch_val_images, batch_new_val_image_paths))
    pool.close()
    pool.join()
    print("Finish generating val dataset")

def run_main():
    """
       這是主函數
    """
    cifar_dataset_dir = os.path.abspath("./data/cifar-10-batches-py")
    new_cifar_dataset_dir = os.path.abspath("./data/cifar10")
    batch_size = 2000
    cifar_preprocess(cifar_dataset_dir, new_cifar_dataset_dir,batch_size)

if __name__ == '__main__':
    run_main()

結果如下：

---

二、訓練階段Mini-VGG的keras實現

2.1 Mini-VGG的網絡架構

由於CIFAR10數據集中所有圖片的分辨率爲32 * 32，VGG16的下采樣率爲32，那麼使用VGG16來實現CIFAR10數據集的分類任務，那麼CIFAR10數據集的圖像在經過VGG16的卷積模塊作用下提取得到特徵維度爲1 * 1 * 1024。那麼這將導致大量特徵丟失，反而不利於圖像分類。因此爲了技能提取得到的特徵又能使得特徵圖不爲1 * 1 * 1024，在本次項目中我們對VGG16的結構進行有所刪減，形成Mini-VGG。

Mini-VGG的網絡架構爲：
第一層：INPUT =>

第二層：CONV => ReLU => BN => CONV => ReLU=>BN =>MAXPOOL => DROPOUT =>

第三層：CONV =>ReLU => BN => CONV =>ReLU =>BN => MAXPOOL => DROPOUT =>

第四層：FC => ReLU => BN => DROPOUT =>

第五層：FC => SOFTMAX

Layer Type	Output Size	Filter Size/Stride
Input Imgae	32 * 32 * 3
Conv	32 * 32 * 32	3 * 3，K=32
ReLU	32 * 32 * 32
BN	32 * 32 * 32
Conv	32 * 32 * 32	3 * 3，K=32
ReLU	32 * 32 * 32
BN	32 * 32 * 32
MaxPool	16 * 16 * 32	2 * 2
Dropout	16 * 16 * 32
Conv	16* 16 * 64	3 * 3，K=64
ReLU	16* 16 * 64
BN	16* 16 * 64
Conv	16* 16 * 64	3 * 3，K=64
ReLU	16* 16 * 64
BN	16* 16 * 64
MaxPool	8* 8 * 64	2 * 2
Dropout	8 * 8 * 64
FC	512
ReLU	512
BN	512
Dropout	512
FC	10
Softmax	10

2.2 Mini-VGG訓練

Mini-VGG的keras實現與利用數據集生成器進行訓練的代碼如下，由於訓練過程中涉及較多參數，爲了在Mini-VGG類代碼編寫過程中指定過多參數，首先實現參數配置類config用來保存訓練過程中所有相關參數，並實現將所有參數保存到本地txt文件函數，方便訓練過後查看每次訓練的相關細節。參數配置類config的定義如下：

# -*- coding: utf-8 -*-
# @Time    : 2020/5/30 16:12
# @Author  : Dai PuWei
# @Email   : [email protected]
# @File    : config.py
# @Software: PyCharm

import os

class config(object):

    default_dict = {
        "dataset_dir": os.path.abspath("./data/cifar10"),
        "checkpoints_dir": os.path.abspath("./checkpoints"),
        "logs_dir": os.path.abspath("./logs"),
        "result_dir": os.path.abspath("./result"),
        "config_dir": os.path.abspath("./config"),
        "input_image_shape": (32, 32, 3),
        "pre_model_path": None,
        "bacth_size": 16,
        "init_learning_rate": 0.01,
        "epoch": 50,
    }

    def __init__(self,**kwargs):
        """
        這是VGG16的初始化函數
        :param cfg: 參數配置類
        """
        # 初始化相關參數
        self.__dict__.update(self.default_dict)
        self.__dict__.update(kwargs)

        # 初始化相關目錄
        if not os.path.exists(self.checkpoints_dir):
            os.mkdir(self.checkpoints_dir)
        if not os.path.exists(self.logs_dir):
            os.mkdir(self.logs_dir)
        if not os.path.exists(self.result_dir):
            os.mkdir(self.result_dir)
        if not os.path.exists(self.config_dir):
            os.mkdir(self.config_dir)

    def save_logs(self, time):
        """
        這是保存模型訓練相關參數的函數
        :param time: 時間
        :return:
        """
        # 創建本次訓練相關目錄
        self.checkpoints_dir = os.path.join(self.checkpoints_dir, time)
        self.logs_dir = os.path.join(self.logs_dir, time)
        self.config_dir = os.path.join(self.config_dir, time)
        self.result_dir = os.path.join(self.result_dir, time)
        if not os.path.exists(self.config_dir):
            os.mkdir(self.config_dir)
        if not os.path.exists(self.checkpoints_dir):
            os.mkdir(self.checkpoints_dir)
        if not os.path.exists(self.logs_dir):
            os.mkdir(self.logs_dir)
        if not os.path.exists(self.result_dir):
            os.mkdir(self.result_dir)

        config_txt_path = os.path.join(self.config_dir, "config.txt")
        with open(config_txt_path, 'a') as f:
            for key, value in self.__dict__.items():
                s = key + ": " + str(value) + "\n"
                f.write(s)

接下來是訓練階段Mini-VGG類的定義如下：

# -*- coding: utf-8 -*-
# @Time    : 2020/5/24 17:11
# @Author  : Dai PuWei
# @Email   : [email protected]
# @File    : mini_vgg_train.py
# @Software: PyCharm

import os
import datetime
import numpy as np
from matplotlib import pyplot as plt

from keras.layers import Input
from keras.layers import Dense
from keras.layers import Dropout
from keras.layers import Flatten
from keras.layers import MaxPooling2D
from keras.layers import Convolution2D
from keras.layers.normalization import BatchNormalization

from keras import Model
from keras.optimizers import Adam
from keras import backend as K

from keras.callbacks import TensorBoard
from keras.callbacks import EarlyStopping
from keras.callbacks import ModelCheckpoint
from keras.callbacks import ReduceLROnPlateau

class mini_VGG(object):

    def __init__(self,cfg):
        """
        這是VGG16的初始化函數
        :param cfg: 參數配置類
        """
        # 初始化相關參數
        self.cfg = cfg

        # 搭建MiNi-VGG，並編譯模型
        self.build_model()
        """
        self.model.compile(optimizer=SGD(lr=self.init_learning_rate, momentum=0.9,
                                         nesterov=True,decay= 0.01 / self.epoch),
                           loss=["categorical_crossentropy"],metrics=["acc"])
        """
        self.model.compile(optimizer=Adam(lr=self.cfg.init_learning_rate),
                           loss=["categorical_crossentropy"], metrics=["acc"])
        if self.cfg.pre_model_path is not None:
            self.model.load_weights(self.cfg.pre_model_path,by_name=True,skip_mismatch=True)
            print("loads model from: ",self.cfg.pre_model_path)

    def build_model(self):
        """
        這是Mini-VGG網絡的搭建函數
        :return:
        """
        # 初始化網絡輸入
        self.image_input = Input(shape=self.cfg.input_image_shape,name="image_input")

        y = Convolution2D(filters=64, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(self.image_input)
        y = BatchNormalization()(y)
        y = Convolution2D(filters=64, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = MaxPooling2D(pool_size=2, strides=2, padding='same')(y)
        y = Dropout(0.25)(y)

        y = Convolution2D(filters=128, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = Convolution2D(filters=128, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = MaxPooling2D(pool_size=2, strides=2, padding='same')(y)

        y = Convolution2D(filters=256, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = Convolution2D(filters=256, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = MaxPooling2D(pool_size=2, strides=2, padding='same')(y)
        y = Dropout(0.25)(y)

        y = Flatten()(y)
        y = Dense(512, activation='relu', kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = Dropout(0.5)(y)
        y = Dense(10, activation='softmax', kernel_initializer='he_normal')(y)

        self.model = Model(self.image_input,y,name="Mini-VGG")
        self.model.summary()

    def train(self,train_datagen,val_datagen,train_iter_num,val_iter_num,init_epoch=0):
        """
        這是VGG16的訓練函數
        :param train_datagen: 訓練數據集生成器
        :param val_datagen: 驗證數據集生成器
        :param train_iter_num: 一個epoch訓練迭代次數
        :param val_iter_num: 一個epoch驗證迭代次數
        :param init_epoch: 初始週期數
        """
        # 初始化相關文件目錄路徑,並保存到日誌文件
        time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        self.cfg.save_logs(time)

        # 初始化回調函數
        tensorboard = TensorBoard(self.cfg.logs_dir,)
        early_stop = EarlyStopping(monitor='val_loss',min_delta=1e-6,verbose=1,patience=10)
        reduce_lr = ReduceLROnPlateau(monitor='val_loss',factor=0.5,verbose=1,patience=2)
        checkpint1 = ModelCheckpoint(filepath=os.path.join(self.cfg.checkpoints_dir,
                                                          'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}-'
                                                          'acc{acc:.3f}-val_acc{val_acc:.3f}.h5'),
                                    monitor='val_loss', save_best_only=True,verbose=1)
        checkpint2 = ModelCheckpoint(filepath=os.path.join(self.cfg.checkpoints_dir,
                                                          'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}-'
                                                          'acc{acc:.3f}-val_acc{val_acc:.3f}.h5'),
                                    monitor='val_acc', save_best_only=True, verbose=1)

        # 訓練模型--第一階段
        history1 = self.model.fit_generator(train_datagen,steps_per_epoch=train_iter_num,
                                 validation_data=val_datagen,validation_steps=val_iter_num,verbose=1,
                                 initial_epoch=init_epoch,epochs=self.cfg.epoch,
                                 callbacks=[tensorboard,checkpint1,checkpint2,early_stop,reduce_lr])
        self.model.save(os.path.join(self.cfg.checkpoints_dir,"stage1-trained-model.h5"))

        # 凍結除最後預測分類的全連接層之外的所有層參數
        for i in range(len(self.model.layers)-1):
            self.model.layers[i].trainable = False

        # 重新設置學習率
        K.set_value(self.model.optimizer.lr,self.cfg.init_learning_rate // 100)

        # 初始化回調函數
        checkpint1 = ModelCheckpoint(filepath=os.path.join(self.cfg.checkpoints_dir,
                                                           'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}-'
                                                           'acc{acc:.3f}-val_acc{val_acc:.3f}.h5'),
                                     monitor='val_loss', save_best_only=True, verbose=1)
        checkpint2 = ModelCheckpoint(filepath=os.path.join(self.cfg.checkpoints_dir,
                                                           'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}-'
                                                           'acc{acc:.3f}-val_acc{val_acc:.3f}.h5'),
                                     monitor='val_acc', save_best_only=True, verbose=1)
        reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, verbose=1, patience=2)
        history2 = self.model.fit_generator(train_datagen, steps_per_epoch=train_iter_num,
                                 validation_data=val_datagen, validation_steps=val_iter_num, verbose=1,
                                 initial_epoch=self.cfg.epoch, epochs=self.cfg.epoch*2,
                                 callbacks=[tensorboard, checkpint1,checkpint2,early_stop,reduce_lr])
        self.model.save(os.path.join(self.cfg.checkpoints_dir, "stage2-trained-model.h5"))

        # 繪製訓練與驗證損失走勢圖
        loss = np.concatenate([history1.history["loss"],history2.history["loss"]])
        val_loss = np.concatenate([history1.history["val_loss"], history2.history["val_loss"]])
        plt.plot(np.arange(0, len(loss)), loss, label="train_loss")
        plt.plot(np.arange(0, len(val_loss)), val_loss, label="val_loss")
        plt.title("Loss on CIFAR-10")
        plt.xlabel("Epoch")
        plt.ylabel("Loss")
        plt.legend()
        plt.grid(True)
        plt.savefig(os.path.join(self.cfg.result_dir,"loss.png"))
        plt.close()

        # 繪製訓練與驗證精度走勢圖
        acc = np.concatenate([history1.history["acc"], history2.history["acc"]])
        val_acc = np.concatenate([history1.history["val_acc"], history2.history["val_acc"]])
        plt.plot(np.arange(0, len(acc)), acc, label="train_acc")
        plt.plot(np.arange(0, len(val_acc)), val_acc, label="val_acc")
        plt.title("Accuracy on CIFAR-10")
        plt.xlabel("Epoch")
        plt.ylabel("Accuracy")
        plt.legend()
        plt.grid(True)
        plt.savefig(os.path.join(self.cfg.result_dir, "accuracy.png"))
        plt.close()

那麼Mini-VGG的訓練腳本如下，在訓練過程中爲了增加模型的魯棒性，對於訓練集和測試集都進行相關的數據增強，包括旋轉、裁剪、水平翻轉、垂直翻轉、亮度變化等。

# -*- coding: utf-8 -*-
# @Time    : 2020/5/24 22:53
# @Author  : Dai PuWei
# @Email   : [email protected]
# @File    : train_mini_vgg.py
# @Software: PyCharm

import os
#os.environ["CUDA_VISIBLE_DEVICES"] = "3"

from config.config import config
from model.mini_vgg_train import mini_VGG
from keras.preprocessing.image import ImageDataGenerator

def run_main():
    """
       這是主函數
    """
    # 初始化參數配置類
    batch_size = 128
    epoch = 50
    cfg = config(epoch = epoch,batch_size=batch_size)

    # 構造訓練集和測試集數據生成器
    train_dataset_dir = os.path.abspath("./data/cifar10/train")
    val_dataset_dir = os.path.abspath("./data/cifar10/val")
    image_data =  ImageDataGenerator(rotation_range=0.2,
                                    width_shift_range=0.05,
                                    height_shift_range=0.05,
                                    shear_range=0.05,
                                    zoom_range=0.05,
                                    horizontal_flip=True,
                                    vertical_flip=True,
                                    rescale= 1.0/255)
    train_datagen = image_data.flow_from_directory(train_dataset_dir,
                                                   class_mode='categorical',
                                                   batch_size = batch_size,
                                                   target_size=(32,32),
                                                   shuffle=True)
    val_datagen = image_data.flow_from_directory(val_dataset_dir,
                                                 class_mode='categorical',
                                                 batch_size=batch_size,
                                                 target_size=(32,32),
                                                 shuffle=True)
    train_iter_num = train_datagen.samples // batch_size
    val_iter_num = val_datagen.samples // batch_size
    if train_datagen.samples % batch_size != 0:
        train_iter_num += 1
    if val_datagen.samples % batch_size != 0:
        val_iter_num += 1

    # 初始化VGG16，並進行測試批量圖像
    mini_vgg = mini_VGG(cfg)
    mini_vgg.train(train_datagen=train_datagen,
                    val_datagen=val_datagen,
                    train_iter_num=train_iter_num,
                    val_iter_num=val_iter_num)

if __name__ == '__main__':
    run_main()

經過100個epoch的訓練之後，Mini-VGG的訓練與驗證損失、訓練與驗證精度的走勢圖如下圖所示。由於在訓練過程中設置了早停回調函數，因此100個epoch的訓練在不到50個epoch訓練週期就結束了。從下損失和精度的走勢圖可以得知，在驗證集上雖然損失在上下波動，但是精度收斂到了85%附近。因此，Mini-VGG在對CIFAR10數據集經過訓練之後，擁有較高的分類性能。

三、測試階段Mini-VGG的keras實現

在訓練結束之後，下一步就是準確評估Mini-VGG在數據集上的性能。測試階段的Mini-VGG類的定義如下：

# -*- coding: utf-8 -*-
# @Time    : 2020/5/24 17:11
# @Author  : Dai PuWei
# @Email   : [email protected]
# @File    : mini_vgg_test.py
# @Software: PyCharm

import os
import numpy as np
from sklearn.metrics import classification_report

from keras import Model
from keras.models import load_model
from keras.layers import Input
from keras.layers import Dense
from keras.layers import Dropout
from keras.layers import Flatten
from keras.layers import MaxPooling2D
from keras.layers import Convolution2D
from keras.layers.normalization import BatchNormalization

class mini_VGG(object):

    _default_dict_ = {
        "input_image_shape": (32,32,3),
        "model_path": os.path.abspath("./data/mini_vgg.h5")
    }

    def __init__(self,**kwargs):
        """
        這是VGG16的初始化函數
        :param kwargs: 參數字典
        """
        # 初始化相關參數
        self.__dict__.update(self._default_dict_)
        self.__dict__.update(kwargs)

        # 加載模型
        try:
            self.model = load_model(self.model_path)
        except:
            self.build_model()      # 搭建MiNi-VGG
            self.model.load_weights(self.model_path,by_name=True,skip_mismatch=True)
        print("loads model from: ",self.model_path)

    def build_model(self):
        """
        這是Mini-VGG網絡的搭建函數
        :return:
        """
        # 初始化網絡輸入
        self.image_input = Input(shape=self.input_image_shape,name="image_input")

        y = Convolution2D(filters=64, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(self.image_input)
        y = BatchNormalization()(y)
        y = Convolution2D(filters=64, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = MaxPooling2D(pool_size=2, strides=2, padding='same')(y)
        y = Dropout(0.25)(y)

        y = Convolution2D(filters=128, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = Convolution2D(filters=128, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = MaxPooling2D(pool_size=2, strides=2, padding='same')(y)

        y = Convolution2D(filters=256, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = Convolution2D(filters=256, kernel_size=3, strides=1, padding='same', activation='relu',
                          kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = MaxPooling2D(pool_size=2, strides=2, padding='same')(y)
        y = Dropout(0.25)(y)

        y = Flatten()(y)
        y = Dense(512, activation='relu', kernel_initializer='he_normal')(y)
        y = BatchNormalization()(y)
        y = Dropout(0.5)(y)
        y = Dense(10, activation='softmax', kernel_initializer='he_normal')(y)

        self.model = Model(self.image_input,y,name="Mini-VGG")
        self.model.summary()

    def eval_generator(self,datagen,iter_num,label_names):
        """
        這是利用數據集生成器對模型進行評估的函數
        :param datagen: 數據集生成器
        :param iter_num: 數據集生成器迭代次數
        :param label_names: 標籤名稱
        :return:
        """
        y_real = []
        y_pred = []
        for i in np.arange(iter_num):
            batch_images,batch_real_labels = datagen.__next__()
            y_real.append(np.argmax(batch_real_labels,axis=-1))
            batch_pred_labels = self.model.predict_on_batch(batch_images)
            y_pred.append(np.argmax(batch_pred_labels,axis=-1))
        y_real = np.concatenate(y_real)
        y_pred = np.concatenate(y_pred)

        return classification_report(y_real,y_pred,target_names=label_names)

在一個數據集上評估Mini-VGG性能的腳本如下：

# -*- coding: utf-8 -*-
# @Time    : 2020/3/24 19:15
# @Author  : Dai PuWei
# @Email   : [email protected]
# @File    : eval_on_dataset.py
# @Software: PyCharm

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"

from model.mini_vgg_test import mini_VGG
from keras.preprocessing.image import ImageDataGenerator

def run_main():
    """
       這是主函數
    """
    # 初始化參數配置類
    image_data = ImageDataGenerator(rotation_range=0.2,
                                    width_shift_range=0.05,
                                    height_shift_range=0.05,
                                    shear_range=0.05,
                                    zoom_range=0.05,
                                    horizontal_flip=True,
                                    vertical_flip=True,
                                    rescale=1.0 / 255)

    # 構造驗證集數據生成器
    #dataset_dir = os.path.join(cfg.dataset_dir, "train")
    dataset_dir = os.path.abspath("./data/cifar10/val")
    image_data = ImageDataGenerator(rescale=1.0 / 255)
    datagen = image_data.flow_from_directory(dataset_dir,
                                             class_mode='categorical',
                                             batch_size = 1,
                                             target_size=(32,32),
                                             shuffle=False)

    # 初始化相關參數
    iter_num = datagen.samples       # 訓練集1個epoch的迭代次數
    label_names = ['airplane','automobile','bird','cat','deer','dog','frog','horse','ship','truck']

    # 初始化VGG16，並進行訓練
    model_path = os.path.abspath("./checkpoints/20200530195246/stage2-trained-model.h5")
    mini_vgg = mini_VGG(model_path=model_path)
    print(mini_vgg.eval_generator(datagen,iter_num,label_names))

if __name__ == '__main__':
    run_main()

Mini-VGG在CIFAR10的驗證集上的評估結果如下：