텐서플로우 Dataset: from_generator 설명

 

from_generator()

사용 이유

1. 메모리 용량의 한계 때문에

- 이미지, 텍스트 데이터 등을 메모리에 올려 놓은 후 작업을 진행할 수 있지만, 수많은 파일이 있을 경우 이를 모두 메모리에 올려놓지 못하기에, generator로부터 1개 또는 배치단위만큼 파일을 불러와야 합니다.

2. 다양한 가공 처리를 위해서

- 간혹 이미지 채널을 기존 RGB 3채널에서 4,5 채널로 늘리는 경우 또는 이미지의 특정 부분을 자른다거나 겹치게 하는 등의 다양한 처리를 generator를 통해서 할 수 있기 때문입니다.

 

필수 사항

1. 모델의 Input shape, Output shape

2. generator: 이미지와 해당 라벨을 추출하는 제너레이터

3. 모델에 넣을 데이터의 자료형과 라벨의 자료형(Int, Float, String 등)

 

 

코드 설명

이미지 로드 예시

 

1. 모델을 정의하고, Input shape와 Output shape를 정합니다.

여기서는 두 개의 RGB 이미지를 겹쳐서 넣는 CNN 모델을 제작해 봅시다.

이미지 사이즈는 (512, 512) 라고 가정합시다.

(512, 512, 3) + (512, 512, 3) => (512, 512, 6)

model = keras.Sequential(
    [
        keras.Input(shape=(512, 512, 6)), # input shape 정의!
        layers.Conv2D(filters=24, kernel_size=(3, 3), padding="same",),
        layers.MaxPooling2D(),
        layers.BatchNormalization(),
        layers.GlobalAveragePooling2D(),
        layers.Dense(2, activation='softmax') # output shape 정의
    ]
)
model.compile(loss="categorical_crossentropy", optimizer="adam")
model.summary()

 

2. 제너레이터를 생성합니다.

# 두 이미지를 겹치게 만드는 generator를 만들어 봅시다!

# 아래의 train_label은 train dataset 정보가 담긴 pd.dataframe 입니다.
# pd.dataframe이 아니더라도, 이미지 파일명이 담긴 리스트 등 또한 가능합니다.

# 우선 이미지0와 이미지1을 담은 리스트 또는 pd.Dataframe을 준비합니다.
files0 = train_label['img0_name']
files1 = train_label['img1_name']

labels = train_label['label'] # 이미지의 라벨을 담은 pd.Dataframe 또는 list

# generator 정의. yield로 끝나는 걸 볼 수 있습니다.
def gen():
    for i in range(len(train_label)):
        file0 = Image.open(data_dir / files[i]) # pillow 라이브러리를 사용하여 이미지를 불러옵니다.
        file1 = Image.open(data_dir / files[i]) # pillow 라이브러리를 사용하여 이미지를 불러옵니다.

        file0 = np.asarray(file0) # 넘파이로 변환합니다.
        file1 = np.asarray(file1) # 넘파이로 변환합니다.
        
        file = np.concatenate((file0, file1), axis=2) # 두 이미지를 겹칩니다.
        file = file.astype(np.int16) # 이미지를 normalize 하지 않고, [0,255] 범위로 표현합니다.
        
        label = np.array(labels[i]) # 라벨 정보를 불러옵니다.
        label = keras.utils.to_categorical(label, 2) # [0,1]의 라벨을 [[1,0], [0,1]]로 바꿉니다.
        label = label.astype(np.int8) # 자료형을 정수로 바꿉니다.
        
        yield (file, label)
        # file의 shape는 (512, 512, 6), 자료형은 np.int16
        # label의 shape는 (2), 자료형은 np.int8 입니다.

 

3. Dataset.from_generator() 정의

# Dataset.from_generator 정의
dataset = tf.data.Dataset.from_generator(gen, (tf.int16, tf.int8), ((512,512,6), (2)))
dataset = dataset.batch(32)

# from_generator()의 인자로,
# 1. generator: 앞서 정의한 generator인 gen 을 넣어줍니다.
# 2. output_types: 제너레이터로부터 나오는 image와 label의 자료형을 입력합니다. (tf.int16, tf.int8)
#     image의 경우 [0, 255] 범위를 갖기에 tf.int16
#     label의 경우 [0, 1]의 범위를 갖기에 tf.int8 입니다.
# 3. output_shapes: 제너레이터로부터 나오는 image와 label의 shape를 입력합니다. ((512,512,6), (2))

# 그리고 dataset의 batch를 지정해주어야 합니다.
# 자신의 컴퓨터 리소스에 맞춰 batch를 지정해줍니다. 

 

4. 모델 훈련하기

model.fit(dataset) # fit()안에 앞서 만든 데이터셋을 넣어주면 끝입니다.
# 읽어주셔서 감사합니다!