본문 바로가기

DeepLearning/OCR_

[ENG_OCR] CRNN_Training 모델 학습 프로세스

반응형

< Model Training Full Process >

아래에서 진행되는 코드는 아래 깃허브를 기본 베이스로 사용하였습니다:)

https://github.com/mvoelk/ssd_detectors

 

mvoelk/ssd_detectors

SSD-based object and text detection with Keras, SSD, DSOD, TextBoxes, SegLink, TextBoxes++, CRNN - mvoelk/ssd_detectors

github.com

 

CRNN_Train.py

1> 필요한 라이브러리 로딩

import numpy as np
import matplotlib.pyplot as plt
import os
import editdistance
import pickle
import time
from keras.optimizers import SGD

from crnn_model import CRNN
from crnn_data import InputGenerator
from crnn_utils import decode
from utils.training import Logger, ModelSnapshot

 

2> 생성된 피클 파일 불러오기

from data_cocotext import GTUtility

# Train
file_name1 = 'gt_util_cocotext_train.pkl'

with open(file_name1, 'rb') as f:
    gt_util_train = pickle.load(f)
 
# Validation 
file_name2 = 'gt_util_cocotext_val.pkl'

with open(file_name2, 'rb') as f:
    gt_util_val = pickle.load(f)

- 피클 파일 생성에 대한 내용은 다른 포스트에서 참고가능

- 현재 피클 파일은 Train을 위한 파일, Validation을 위한 파일 이렇게 2개가 존재한다

- 저장 된 피클파일을 로드해서 새롭게 변수에 담아준다.

 

 

3> 모델 인풋파라미터에 대한 정의

from crnn_utils import alphabet87 as alphabet
# len(alphabet) = 87

input_width = 256
input_height = 32
batch_size = 128

input_shape = (input_width, input_height, 1)

- 모델에 인풋하기 위해서 input_shape을 정의해준다.

 

 

4> Fine-Tunning을 하기위해 동결 할 레이어층 

freeze = ['conv1_1',
          'conv2_1',
          'conv3_1', 'conv3_2', 
          'conv4_1',
          'conv5_1',
          #'conv6_1',
          #'lstm1',
          #'lstm2'
         ]

- 위에서 선언한 함수는 가중치가 동결되어 학습되는 동안 유지 되도록 한다

- 그외 위에서 선언되지 않는 레이어층은 학습하는동안 가중치가 수정된다.

 

 

5> 모델이 저장 될 버전 이름 정의하기

experiment = 'crnn_lstm_cocotext_pretrained_v10'

- 모델을 학습 할 때마다 새롭게 생기는 log과 history를 저장하기 위해 이름을 정의한다

 

 

6> 인풋을 위한 제너레이터 생성

max_string_len = model_pred.output_shape[1]

from crnn_data import InputGenerator
gen_train = InputGenerator(gt_util_train, batch_size, alphabet, input_shape[:2], 
                           grayscale=True, max_string_len=max_string_len)
gen_val = InputGenerator(gt_util_val, batch_size, alphabet, input_shape[:2], 
                         grayscale=True, max_string_len=max_string_len)

- 아까 변수에 정의한 피클파일을 학습시키기 위한 제너레이터 형태로 만들어주기 위해 InputGenerator를 사용한다

- InputGenerator에서 중간에 crop_word라는 함수를 사용해서 이미지에서 텍스트부분만 추출한다.

 

 

7> 이미 학습 되어있는 가중치를 모델에 로딩

model.load_weights('./checkpoints/201806162129_crnn_lstm_synthtext/weights.300000.h5')

- 이미 학습되어 저장된 가중치를 활용하여 Transfer Learning를한다.

 

 

8> 학습되는 동안의 과정을 저장

checkdir = './checkpoints/' + time.strftime('%Y%m%d%H%M') + '_' + experiment
if not os.path.exists(checkdir):
    os.makedirs(checkdir)

with open(checkdir+'/source.py','wb') as f:
    source = ''.join(['# In[%i]\n%s\n\n' % (i, In[i]) for i in range(len(In))])
    f.write(source.encode())

- 학습하는 동안 log와 history를 저장하기 위해 폴더를 생성하고 파일을 만든다

 

 

9> 옵티마이저 설정

optimizer = SGD(lr=0.0001, decay=1e-6, momentum=0.9, nesterov=True, clipnorm=5)

- 경사하강법(SGD)을 사용하고 learning rate = 0.0001 로 지정

- 러닝레이트를 0.01~0.0001까지 다양하게 변경하면서 학습해본 결과 0.0001의 성능이 가장 좋다

 >> 러닝레이트가 0.01일 경우에는 오버피팅이 빠르게 발생했고, 0.0001보다 더 작을 경우는 큰 변화가 없었다. 

 

 

10> 위에서 선언한 레이어를 가중치 동결

for layer in model.layers:
    layer.trainable = not layer.name in freeze

- conv1_1 부터 conv5_1까지 동결하고, 이후에 conv넷의 마지막층과 RNN층은 동결을 해지한다.

 

 

11> 모델 컴파일 하기

model.compile(loss={'ctc': lambda y_true, y_pred: y_pred}, optimizer=optimizer)

- 모델을 학습시키기 전에 컴파일을 한다. 

- loss는 모델을 구성할 때 정의 한 ctc loss를 사용한다.

 

 

12> 모델 학습시키기

hist = model.fit_generator(generator=gen_train.generate(), # batch_size here?
                    	   steps_per_epoch=gt_util_train.num_objects // batch_size,
                    	   epochs=50,
                    	   validation_data=gen_val.generate(), # batch_size here?
                    	   validation_steps=gt_util_val.num_objects // batch_size,
                    	   callbacks=[
         	ModelCheckpoint(checkdir+'/weights.{epoch:03d}.h5', verbose=1, save_weights_only=True),
         	Logger(checkdir),
                    			],
                    		initial_epoch=0)
반응형