[Tensorflow] VGG16 모델을 이용하여 CNN 이미지 분류기 학습하기

1. Concept

지난 포스팅에 이어, 이번 포스팅에서는 특정한 객체를 집중적으로 분류하기 위해 사전 학습된 신경망 모델을 기반으로 가장 기초적인 방법을 통해 미세 학습 (Find-Tuning)을 구현해 보록 하겠습니다. 

 

ImageNet으로 학습된 VGG16 모델을 기반으로, Kaggle에서 제공되는 고양이 강아지 분류 데이터를 활용하겠습니다. 데이터는 200mb 정도이며 Kaggle 원문 링크 또는 구글 드라이브에서 받으실 수 있습니다.

 

 

그림 1. Cat and Dog 분류를 위한 신경망 모델 (원문 링크 글 하단 참조)

 

구현 환경 : Windows 10 / Conda / Python 3.8 / Tensorflow 2.2 / CUDA 10.2

이전 포스팅을 참고하시면 기초적인 텐서플로우 구현에 도움이 됩니다

 

1. [텐서플로우] 아나콘다 가상 환경에서 텐서플로우 설치하기
2. [텐서플로우] 텐서플로우에서 사전 학습된 VGG16 모델 불러오기
3. [텐서플로우] 사전 학습된 VGG16 모델로 이미지 분류하기

 

2. Python Code

1. 소스코드 다운로드 : github.com/DEEPI-LAB/python-TensorFlow-Tutorials.git

DEEPI-LAB/python-TensorFlow-Tutorials

git clone https://github.com/DEEPI-LAB/python-TensorFlow-Tutorials.git

2.  사전학습된 VGG16 모델 불러오기

from tensorflow.python.keras.applications.vgg16 import VGG16
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator
from tensorflow.python.keras.layers import Conv2D, MaxPooling2D, Flatten,
from tensorflow.python.keras.layers import Dense, Dropout, Input
from tensorflow.python.keras.models import Model

# 사전 학습된 모델 불러오기
input_tensor = Input(shape=(150,150,3))
model = VGG16(weights='imagenet', include_top=False, input_tensor = input_tensor)

 

VGG16 MODEL 함수

 

weights : 가중치 모델 지정 ( None : 초기화된 가중치, 'imagenet' : 사전 학습된 가중치 )

include_top : 신경망 FC 층 존재 유무 ( False : 제거 / True : 유지 )

input_tensor : 입력 텐서 크기 ( Input(shape = (w, h, ch))

 

 

사전 학습된 VGG16 모델을 불러옵니다. Find-Tuning으로 새로운 분류기를 학습하기 위해 기존의 FC (Fully-Connected Layer)를 제거하고 입력되는 이미지의 크기 input_tensor를 지정합니다.

 

input_tensor는 사용되는 분류기 특성이나 학습 이미지의 크기 등을 고려하여 설정하면 됩니다. 하지만 include_top이 True일 경우, 입력 이미지의 크기는 사전 학습된 모델과 같이 224x224x3으로 고정되기 때문에 input_tensor로 변경할 수 없습니다.

 

3. VGG 신경망 모델 디자인

# 모델 Layer 데이터화
layer_dict = dict([(layer.name, layer) for layer in model.layers])

# Layer 추가
x = layer_dict['block5_pool'].output
# Cov2D Layer +
x = Conv2D(filters = 64, kernel_size=(3, 3), activation='relu')(x)
# MaxPooling2D Layer +
x = MaxPooling2D(pool_size=(2, 2))(x)
# Flatten Layer +
x = Flatten()(x)
# FC Layer +
x = Dense(2048, activation='relu')(x)
x = Dropout(0.5)(x)
x = Dense(1024, activation='relu')(x)
x = Dropout(0.5)(x)
x = Dense(2, activation='softmax')(x)

# new model 정의
new_model = Model(inputs = model.input, outputs = x)

 

신경망 기본 Layer 함수

 

Conv2D( filters = 채널 수, kernel_size = CNN 크기, activation = 활성화 함수 )

MaxPooling2D(pool_size = 폴링 사이즈)

Flatten() CNN -FC 연결을 위한 함수

Dense( 노드 수, activation = 활성화 함수 ('relu', 'sigmoid') )

Dropout( 계수 (0 ~ 1.0) )

 

기본적인 신경망 모델 디자인 방법은 단순합니다. VGG16 모델의 OUTPUT을 분기점으로 하여 새로운 모델 X를 생성한 뒤, 이어 나가면 됩니다. 이후, Model 함수를 통해 입력과 출력을 융합하여 new model를 정의합니다.

 

예제에서는 64 채널을 가지는 3x3 크기의 CNN과 MaxPooling 레이어를 추가 한 뒤, 2048 - 1024 - 2의 FC를 추가하였습니다. 마지막 출력층은 고양이와 강아지를 분류하는 신경망이기 때문에 2개의 출력, softmax 함수로 구성하였습니다.

 

 

4.  새롭게 디자인된 VGG 모델 

# CNN Pre-trained 가중치를 그대로 사용할때
for layer in new_model.layers[:19] :
    layer.trainable = False

new_model.summary()

# 컴파일 옵션
new_model.compile(loss='sparse_categorical_crossentropy',
                     optimizer='adam',
                     metrics=['accuracy'])

 

사전 학습된 VGG16를 Find-Tuning 한다는 것은 전체 네트워크의 가중치를 새롭게 학습하는 것이 아닌 새롭게 추가된 레이어나 말단 부분만을 학습하는 것입니다. 이를 위해 layer.trainable = False를 선언합니다. 새롭게 추가된 층 이외의 VGG16의 19번째 레이어까지 학습하지 않고 고정하였습니다.

 

 

5. 학습 데이터 설정

# 테스트 데이터 (고양이/강아지)
train_dir ='train'
test_dir = 'test'

# 폴더에 따라 자동 분류
train_image_generator = ImageDataGenerator(rescale=1./255)
test_image_generator = ImageDataGenerator(rescale=1./255)

# 데이터 구조 생성
train_data_gen = train_image_generator.flow_from_directory(batch_size=16,
                                                           directory=train_dir,
                                                           shuffle=True,
                                                           target_size=(150, 150),
                                                           class_mode='binary')

test_data_gen = test_image_generator.flow_from_directory(batch_size=16,
                                                         directory=test_dir,
                                                         target_size=(150, 150),
                                                         class_mode='binary')

 

학습을 위한 데이터 전처리는 다양한 방법으로 가능합니다. 본 포스팅에서는 이진 분류기에 적용하기 쉬운 ImageData Generator를 활용하였습니다.

 

모든 이미지는 신경망 학습 이전에 1./255를 통해 float 형태로 전 처리됩니다.

 

 

flow_from_directory 함수

 

batch_size (학습을 위한 배치 사이즈)

directory (학습 데이터 위치)

shuffle (데이터 셔플링)

target_size (입력 이미지 크기)

class_mode (라벨)

 

 

6. 신경망 모델 학습

# 모델 학습
history = new_model.fit(train_data_gen, epochs=5,
                        validation_data=test_data_gen)

new_model.save("newVGG16.h5")                        
                        
# 최종 결과 리포트
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']
loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(len(acc))

from matplotlib import pyplot as plt

plt.plot(epochs, acc, 'r', label='Training acc')
plt.plot(epochs, val_acc, 'b', label='testing acc')
plt.title('Training and testing accuracy')
plt.legend()
plt.figure()

plt.plot(epochs, loss, 'r', label='Training loss')
plt.plot(epochs, val_loss, 'b', label='testing loss')
plt.title('Training and testing loss')
plt.legend()

plt.show()

# 저장 모델 불러오기
from keras.models import load_model

new_model = load_model("newVGG16.h5")

 

fit을 통해 신경망을 학습합니다. epochs는 반복 학습 횟수이며 학습 데이터와 검증 데이터만 지정해주면 간단한 신경망 학습이 시작되며 콘솔 창에서 실시간 학습 현황이 출력됩니다.

 

모든 학습이 종료되면 matplotlib를 통해 결과를 그림 2와 같이 그래프화 할 수 있습니다.

 

그림 2. 최종 학습 결과 그래프

 

Your Best AI Partner DEEP.I
AI 바우처 공급 기업
객체 추적 및 행동 분석 솔루션 | 제조 생산품 품질 검사 솔루션 | AI 엣지 컴퓨팅 시스템 개발
인공지능 프로젝트 개발 외주 및 상담
E-mail: contact@deep-i.ai
Site: www.deep-i.ai

딥아이 DEEP.I | AI 기반 지능형 기업 솔루션