Transfer Learning을 이용한 이미지 분류 실습

강아지와 고양이를 분류하는 실습을 Transfer Learning을 이용하여 해봅시다.

 

실습은 Colab으로 진행했습니다.

 

일단 파일을 다운로드 받습니다.

!wget --no-check-certificate \
    https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip \
    -O ./cats_and_dogs_filtered.zip

 

import os
import zipfile
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing.image import ImageDataGenerator

 

다운로드 받은 파일의 압축을 풉니다.

zip_ref = zipfile.ZipFile('/content/cats_and_dogs_filtered.zip')
zip_ref.extractall('./')
zip_ref.close()

데이터셋의 경로를 지정합니다.

train_dir = '/content/cats_and_dogs_filtered/train'
validation_dir = '/content/cats_and_dogs_filtered/validation'

 

미리 잘 학습된 모델을 MobileNetV2를 이용해 불러옵니다.

- MobileNetV2란 모바일이나, 임베디드에서도 실시간을 작동할 수 있게 모델이 경량화 되면서도, 

  정확도 또한 많이 떨어지지 않게하여, 속도와 정확도 사이의 트레이드 오프 문제를 어느정도 해결한 네트워크 입니다.

 

# 우리가 만들려는 모델의 인풋 이미지는 128, 128, 3 으로 한다
IMG_SHAPE = (128, 128, 3)

# 이미 만들어져서 학습이 완료된 모델을 가져오되, head model을 제외한 base model만 가져온다.
base_model =  tf.keras.applications.MobileNetV2(
    input_shape = IMG_SHAPE, include_top = False
)

 

summary() 함수를 통해 정보를 보면 매우 많은 레이어들과 weight로 이루어져 있는 것을 볼 수 있습니다.

base_model.summary()

# base_model의 레이어 수
>>> len(base_model.layers)
154

 

이제 베이스모델 부분은 이미 이미지의 특징을 잡아내도록

잘 학습된 상태이기 때문에 학습이 되지 않도록 설정합니다.

base_model.trainable = False

 

Head Model을 만들어 봅시다.

# base_model의 아웃풋을 받아 head_model의 인풋으로 받습니다.
head_model = base_model.output

# 이제 ANN Layer를 생성하는데
# 추가하고자 하는 Layer()(head_model) 
# 와 같은 형식으로 합니다.

# flatten 되도록 합니다.
head_model = tf.keras.layers.GlobalAveragePooling2D()(head_model)

head_model = tf.keras.layers.Dense(128, 'relu')(head_model)

# output layer
head_model = tf.keras.layers.Dense(1, 'sigmoid')(head_model)

 

Base Model과 Head Model을 하나의 Model로 연결합니다.

# 베이스 모델과 헤드모델을 합쳐서, 하나의 모델이라는 변수로 만들어 준다.
model = tf.keras.models.Model(inputs = base_model.input, outputs = head_model)

 

model을 compile합니다.

model.compile(tf.keras.optimizers.RMSprop(0.0001), loss = 'binary_crossentropy', metrics = ['accuracy'])​

 

 

트레인 이미지의 다양성을 위해 이미지 증강하도록 하고,

폴더에 있는 이미지들을 트레인 데이터와 밸리데이션 데이터로 나눈다.

train_datagen = ImageDataGenerator(rescale = 1/255.0,
                                   rotation_range=20, 
                                   width_shift_range = 0.2,
                                   height_shift_range = 0.2,
                                   zoom_range = 0.2,
                                   horizontal_flip = True)

validation_datagen = ImageDataGenerator(rescale = 1/255.0)

train_generator = train_datagen.flow_from_directory(train_dir, target_size=(128,128),
                                  class_mode = 'binary',
                                  batch_size = 128)
                                  
validation_generator = validation_datagen.flow_from_directory(validation_dir, target_size = (128,128),
                                                              class_mode = 'binary',
                                                              batch_size = 12)

모델을 학습시킨다.

높은 validation_accuracy가 나오는것을 볼 수 있다.

epoch_history = model.fit(train_generator, epochs = 10, validation_data = (validation_generator))