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

程式設計

10 留言

  1. 台中阿任

    感謝大大不吝分享
    最近在嘗試使用 kaggle 的
    https://www.kaggle.com/c/humpback-whale-identification/

    對於kaggle有些問題
    1.kaggle排名是看CSV…但是別人的結果CSV很容易取得,貼上去一下就衝到一半以上了…這是不是怪怪的? 還是我搞錯了什麼?

    2.kaggle有蠻多人分享自己的 kernels…有些從頭跑到尾的notebook我看得懂.但是例如:
    https://www.kaggle.com/ateplyuk/resnext50-sz448-resnext50-sz224-lb-0-684
    我就看不懂他到底怎麼分層與訓練
    要如何實作別人的kernels呢?

    3.在鯨魚判斷這題目,是不是用CNN VGG16就不好呢? 聽我的老師說相當有難度

    感謝大大解答!!!

  2. 台中阿任

    大大好! 不好意思 我遇到一個問題
    就是我看您的程式的時候
    不知道 label 設定在哪邊
    就我的認知
    label 對應都是設定在 CSV 中
    但是我都沒看到有讀取CSV的地方
    不知道大大可否解答!

    我是機器學習的新手
    感謝大大!

  3. Ed90595

    大大您好!我也是機器學習的新手,近幾天按照大大步驟做確實有做出結果,非常感謝!!!
    但不知道怎麼使用這行:
    python3 predict_resnet50.py train/dog.1000[0-2].jpg train/cat.1000[0-2].jpg
    來一次做 dog/cat.10000~10002.jpg 的作業。
    我的環境使用Anaconda去跑,底下python版本為3.6.8。
    沒辦法使用python3這個指令,而是使用: python predict_resnet50.py train/dog.10000.jpg 一次跑一張圖OK。
    但如果要如大大一次使用多個照片的語法那樣下去跑會跑出:
    FileNotFoundError: [Errno 2] No such file or directory: ‘train/dog.1000[0-2].jpg’
    這個錯誤QQ 拜託大大求解了~~~

    • draguitar

      python3 > python(改用這個)
      FileNotFoundError: [Errno 2] No such file or directory: ‘train/dog.1000[0-2].jpg’ >>>>> 找不到檔案,路徑設錯

    • 王小凱

      # 從參數讀取圖檔路徑
      files = sys.argv[1:]
      ——————————-以上取消,替換以下方式
      import glob
      files= sorted(glob.glob(‘./sample/’ + ‘*.jpg’))
      #只讀取資料夾內 jpg的檔案
      print(files)
      ——————————————————
      我今天試了一下,跑不出來,所以就改成這樣子了
      直接執行,會把資料夾內的jpg檔案抓出來做成list給他吃

  4. Ben

    你好,感謝你的分享,照著你的步驟做確實可以做出辨識的功能,那如果我今天想要改成辨識數字0跟1該注意什麼

  5. handsome hou

    您好
    很感謝您的教程
    不好意思有個問題請問
    我改變資料夾的輸入
    也有成功讀取到檔案
    Found 4000 images belonging to 2 classes.
    Found 0 images belonging to 0 classes.

    train_batches.next()[0].shape
    #(8, 224, 224, 3)

    但是
    訓練會跳出以下error
    想請問如何解決?
    ValueError: Empty training data.

    • Chiang

      有沒有把資料分成 train 跟 valid放置 看起來像是只抓到其中一個資料夾

  6. A唉

    我好奇的是為何你的凍結層數設在第二層?而且這樣的訓練結果不會過凝合嗎?

  7. Forest

    感謝,你的網路給了我很大的幫助

Comments are Closed