[Tensorflow] 학습을 위한 대용량 훈련 데이터 처리 Data Generator 클래스 만들기

Concept

신경망 모델을 학습하기 위해 입력 데이터를 정리하고 전 처리하는 일은 중요하지만 언제나 귀찮은 일입니다. 특히, 이미지 데이터를 학습하기 위해 텐서 플로우에서 ImageGenerator를 설정하는 것은 소규모 데이터에서는 정리가 쉽지만, 대용량 대규모 데이터를 규격에 맞게 정리하는것은 쉽지 않죠.

 

이번 포스팅에서는 ImageGenerator.flow_ 함수를 직접 class로 만들어 커스터마이징이 쉽고 저장이 용이한 학습 데이터 구축 방법을 구현해보도록 하겠습니다.

 

SourceCode

keras.io/api/preprocessing/image/

Keras documentation: Image data preprocessing

1. Keras - ImageDataGenerator.flow_from_dataframe()

 

MNIST나 CIFAR과 같이 저용량 데이터는 구지 이미지 형식으로 저장할 필요 없이 csv, mat, pkl 파일로 묶어두고 그때그때 사용하면 편하지만, 이미지 크기가 100x100 이상만 넘어도 수만 장의 이미지를 메모리에 담는 것은 쉽지 않습니다. Dataframe으로 불러온 데이터를 데이터화하는 것은 대용량 데이터셋에 적합하지 않습니다.

 

2. Keras - ImageDataGenerator.flow_from_directory()

 

데이터를 메모리에 불러오기 전, 데이터가 있는 폴더의 주소로 이미지의 클래스를 구분하여 정렬해줍니다. 하지만, 지정된 폴더 이내 sub 폴더 형식으로 데이터가 분류되어 있어야 합니다. 한 번 정리하면 편하지만, 그때 그때 수십만 장의 이미지를 일일이 분류하기는 쉽지 않습니다.

 

3. Keras - Sequence 를 이용한 Generator 클래스 생성

 

예제로 설명되는 Generator는 다음 전처리 특성을 가지게 구현하였습니다.

• Batch 사이즈 지정
• 입력 이미지 리사이징
• 입력 이미지 정규화
• 데이터와 라벨 일대일 대응

from keras.utils import to_categorical, Sequence

class DataGenerator(Sequence):
    def __init__(self, path, list_IDs, labels, 
    batch_size, img_size, img_channel, num_classes):
     
        # 데이터셋 경로
        self.path = path
        # 데이터 이미지 개별 주소 [ DataFrame 형식 (image 주소, image 클래스) ]
        self.list_IDs = list_IDs
        # 데이터 라벨 리스트 [ DataFrame 형식 (image 주소, image 클래스) ]
        self.labels = labels
        # 학습 Batch 사이즈
        self.batch_size = batch_size
        # 이미지 리사이징 사이즈
        self.img_size = img_size
        # 이미지 채널 [RGB or Gray]
        self.img_channel = img_channel
        # 데이터 라벨의 클래스 수
        self.num_classes = num_classes
        # 전체 데이터 수
        self.indexes = np.arange(len(self.list_IDs))
   
    def __len__(self):
        len_ = int(len(self.list_IDs)/self.batch_size)
        if len_*self.batch_size < len(self.list_IDs):
            len_ += 1
        return len_
    
    def __getitem__(self, index):
        indexes = self.indexes[index*self.batch_size:(index+1)*self.batch_size]
        list_IDs_temp = [self.list_IDs[k] for k in indexes]
        X, y = self.__data_generation(list_IDs_temp)
        return X, y
            
    def __data_generation(self, list_IDs_temp):
        X = np.zeros((self.batch_size, self.img_size, self.img_size, self.img_channel))
        y = np.zeros((self.batch_size, self.num_classes), dtype=int)
        for i, ID in enumerate(list_IDs_temp):
            img = cv2.imread(self.path+ID)
            img = cv2.resize(img, (self.img_size, self.img_size))
            X[i, ] = img/255
            y[i, ] = to_categorical(self.labels[i], num_classes=self.num_classes)
        return X, y
    

 

클래스 생성 뒤, __init__ 를 통해 입력 파라미터를 초기화해줍니다. 입력되는 값은 총 7개입니다. __data_generation 함수에서 데이터를 전 처리한 뒤, __getitem 함수에서 이를 배치 형식으로 정렬하여 출력해줍니다. 기존 flow_from_directory()와 동일하지만, 서브 폴더가 필요하지 않고 라벨과 클래스를 저장된 데이터프레임 파일에 맞게 자동으로 할당해주는 장점이 있습니다.

Application

train.csv

0.58MB

import pandas as pd

# 이미지 주소 및 클래스 라벨 파일 불러오기
train_labels = pd.read_csv('train.csv')

# 라벨 정보 전처리
# 전체 클래스 수
clss_num = len(train_labels['labels'].unique())
# 클래스 -> 숫자로 변환 (카테고리 형식의 클래스를 원 핫 인코딩)
labels_dict = dict(zip(train_labels['labels'].unique(), range(clss_num)))
train_labels = train_labels.replace({"labels": labels_dict})

tartget_size = 150
img_ch = 3
num_class = 12
batch_size = 32

train_generator = DataGenerator('train_images/', train_labels['image'],
                                train_labels['labels'],
                                batch_size, tartget_size,
                                img_ch, num_class)
 
# 학습
# history = model.fit_generator(train_generator, epochs=1)

 

현재 가공된 데이터셋의 형식에 따라 이 방법이 편할수도 있지만, 불편할 수도 있습니다. 데이터셋의 이름이 클래스 별로 구분되어있거나 하위 폴더가 클래스를 나타내도록 구축되어있다면 굳이, 생성기 클래스를 만들 필요는 없습니다. 자신의 학습 모델 데이터에 맞게 활용하시기 바랍니다.

 

 

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

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