Keras 以 ResNet-50 預訓練模型建立狗與貓辨識程式

這裡示範在 Keras 架構下以 ResNet-50 預訓練模型為基礎，建立可用來辨識狗與貓的 AI 程式。

在 Keras 的部落格中示範了使用 VGG16 模型建立狗與貓的辨識程式，準確率大約為 94%，而這裡則是改用 ResNet50 模型為基礎，並將輸入影像尺寸提高為 224×224，加上大量的 data augmentation，結果可讓辨識的準確率達到 99%。

本篇文章所使用的 Keras 是 TensorFlow 1.8.0 所內建的版本，不同版本可能會有一些差異。

準備資料

從 Kaggle 的網站上下載 train.zip 壓縮檔，然後解壓縮：

# 解壓縮 train.zip

unzip train.zip

建立好目錄結構：

# 建立目錄結構

mkdir -p sample/train/cats

mkdir -p sample/train/dogs

mkdir -p sample/valid/cats

mkdir -p sample/valid/dogs

選擇貓與狗各前 1000 張照片作為訓練資料集，另各取 400 張作為測試資料集：

# 複製圖片

cd train

cp cat.?.jpg cat.??.jpg cat.???.jpg ../sample/train/cats/

cp dog.?.jpg dog.??.jpg dog.???.jpg ../sample/train/dogs/

cp cat.1[0-3]??.jpg ../sample/valid/cats/

cp dog.1[0-3]??.jpg ../sample/valid/dogs/

建立與訓練模型

以下是以 ResNet-50 預訓練模型為基礎，建立與訓練狗與貓辨識模型的程式碼：

from tensorflow.python.keras import backend as K

from tensorflow.python.keras.models import Model

from tensorflow.python.keras.layers import Flatten, Dense, Dropout

from tensorflow.python.keras.applications.resnet50 import ResNet50

from tensorflow.python.keras.optimizers import Adam

from tensorflow.python.keras.preprocessing.image import ImageDataGenerator

# 資料路徑

DATASET_PATH = 'sample'

# 影像大小

IMAGE_SIZE = (224, 224)

# 影像類別數

NUM_CLASSES = 2

# 若 GPU 記憶體不足，可調降 batch size 或凍結更多層網路

BATCH_SIZE = 8

# 凍結網路層數

FREEZE_LAYERS = 2

# Epoch 數

NUM_EPOCHS = 20

# 模型輸出儲存的檔案

WEIGHTS_FINAL = 'model-resnet50-final.h5'

# 透過 data augmentation 產生訓練與驗證用的影像資料

train_datagen = ImageDataGenerator(rotation_range=40,

width_shift_range=0.2,

height_shift_range=0.2,

shear_range=0.2,

zoom_range=0.2,

channel_shift_range=10,

horizontal_flip=True,

fill_mode='nearest')

train_batches = train_datagen.flow_from_directory(DATASET_PATH + '/train',

target_size=IMAGE_SIZE,

interpolation='bicubic',

class_mode='categorical',

shuffle=True,

batch_size=BATCH_SIZE)

valid_datagen = ImageDataGenerator()

valid_batches = valid_datagen.flow_from_directory(DATASET_PATH + '/valid',

target_size=IMAGE_SIZE,

interpolation='bicubic',

class_mode='categorical',

shuffle=False,

batch_size=BATCH_SIZE)

# 輸出各類別的索引值

for cls, idx in train_batches.class_indices.items():

print('Class #{} = {}'.format(idx, cls))

# 以訓練好的 ResNet50 為基礎來建立模型，

# 捨棄 ResNet50 頂層的 fully connected layers

net = ResNet50(include_top=False, weights='imagenet', input_tensor=None,

input_shape=(IMAGE_SIZE[0],IMAGE_SIZE[1],3))

x = net.output

x = Flatten()(x)

# 增加 DropOut layer

x = Dropout(0.5)(x)

# 增加 Dense layer，以 softmax 產生個類別的機率值

output_layer = Dense(NUM_CLASSES, activation='softmax', name='softmax')(x)

# 設定凍結與要進行訓練的網路層

net_final = Model(inputs=net.input, outputs=output_layer)

for layer in net_final.layers[:FREEZE_LAYERS]:

layer.trainable = False

for layer in net_final.layers[FREEZE_LAYERS:]:

layer.trainable = True

# 使用 Adam optimizer，以較低的 learning rate 進行 fine-tuning

net_final.compile(optimizer=Adam(lr=1e-5),

loss='categorical_crossentropy', metrics=['accuracy'])

# 輸出整個網路結構

print(net_final.summary())

# 訓練模型

net_final.fit_generator(train_batches,

steps_per_epoch = train_batches.samples // BATCH_SIZE,

validation_data = valid_batches,

validation_steps = valid_batches.samples // BATCH_SIZE,

epochs = NUM_EPOCHS)

# 儲存訓練好的模型

net_final.save(WEIGHTS_FINAL)

將這段 Python 程式碼儲存為 train_resnet50.py，然後執行它即可進行模型的訓練：

python3 train_resnet50.py

模型訓練完成之後，就會儲存於 model-resnet50-final.h5 這個檔案中。

辨識狗與貓

將模型訓練好之後，就可以使用以下的 Python 指令稿打造自己的狗與貓辨識程式了：

from tensorflow.python.keras import backend as K

from tensorflow.python.keras.models import load_model

from tensorflow.python.keras.preprocessing import image

import sys

import numpy as np

# 從參數讀取圖檔路徑

files = sys.argv[1:]

# 載入訓練好的模型

net = load_model('model-resnet50-final.h5')

cls_list = ['cats', 'dogs']

# 辨識每一張圖

for f in files:

img = image.load_img(f, target_size=(224, 224))

if img is None:

continue

x = image.img_to_array(img)

x = np.expand_dims(x, axis = 0)

pred = net.predict(x)[0]

top_inds = pred.argsort()[::-1][:5]

print(f)

for i in top_inds:

print(' {:.3f} {}'.format(pred[i], cls_list[i]))

將這段 Python 程式碼儲存為 predict_resnet50.py，然後就可以拿來辨識圖片了。我們可以直接使用 Kaggle 所提供的圖片進行測試：

# 辨識狗與貓

python3 predict_resnet50.py train/dog.1000[0-2].jpg train/cat.1000[0-2].jpg

train/dog.10000.jpg

1.000 dogs

0.000 cats

train/dog.10001.jpg

1.000 dogs

0.000 cats

train/dog.10002.jpg

1.000 dogs

0.000 cats

train/cat.10000.jpg

0.998 cats

0.002 dogs

train/cat.10001.jpg

1.000 cats

0.000 dogs

train/cat.10002.jpg

1.000 cats

0.000 dogs

也可以拿自己的圖片放進去進行辨識：

小狗圖片

# 辨識自己的圖片

python3 predict_resnet50.py pug-20180918-01.jpg

pug-20180918-01.jpg

0.999 dogs

0.001 cats

參考資料：JK Jung’s blog