[Tensorflow] 대용량 데이터 학습 방법 (Tensorflow 2)

안녕하세요, 오늘은 Tensorflow 2에서 대용량 데이터를 학습하는 방법에 대해서 알아보겠습니다.

 

먼저, MNIST 데이터 셋(28, 28)을 Resize(512, 512) 해서 정해진 경로에 저장하는 코드는 다음과 같습니다.

 

import tensorflow as tf
from PIL import Image
import os
import cv2

# MNIST 데이터셋 로드
mnist = tf.keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = mnist.load_data()

# 디렉토리 생성
os.makedirs('/mnist_png', exist_ok=True)

# (28, 28) => (512, 512) Resize 후 이미지 저장
def save_image(image, label, index):
    os.makedirs(f'/mnist_png/{label}', exist_ok=True)
    image_path = f'/mnist_png/{label}/{index}.{format}'
    image = cv2.resize(image, (512, 512))
    cv2.imwrite(image_path, image)

for i in range(len(train_images)): 
    image, label = train_images[i], train_labels[i]
    save_image(image, label, i)

 

위의 코드를 통해 mnist_png 디렉토리 내 (512, 512) 이미지들이 저장된 것을 확인할 수 있습니다.

 

해당 대용량 데이터를 학습하기 위한 코드는 다음과 같습니다.

 

import tensorflow as tf
from tensorflow.keras.utils import to_categorical
from sklearn.model_selection import train_test_split
import os
import glob
import numpy as np

def read_image(mnist_dir=None):
  mnist_img_list = glob.glob('./mnist_png/**/*.png', recursive=True)
  mnist_img_array = np.array(mnist_img_list)
  label_array = np.array([os.path.dirname(os.path.basename(img_path)) for img_path in mnist_img_list])

  np.random.seed(200)
  np.random.shuffle(mnist_img_array)
  np.random.seed(200)
  np.random.shuffle(label_array)

  return mnist_img_array, to_categorical(label_array, 10)

def fn(mnist_img_path, label):

  tmp = tf.io.read_file(mnist_img_path)
  decode_img = tf.image.decode_png(tmp, channels=1)
  decode_img = tf.cast(decode_img, tf.float32)

  return decode_img, tf.cast(label, tf.float32)


x_dataset, y_dataset = read_image()

X_train, X_test, y_train, y_test = train_test_split(x_dataset, y_dataset, test_size=0.2, random_state=42, stratify=y_dataset)

train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train)) \ 
                .map(fn, num_parallel_calls = tf.data.experimental.AUTOTUNE) \ 
                .shuffle(buffer_size=256) \ 
                .batch(4) \ 
                .prefetch(tf.data.experimental.AUTOTUNE)

validation_dataset = tf.data.Dataset.from_tensor_slices((X_test, y_test)) \
                     .map(fn, num_parallel_calls=tf.data.experimental.AUTOTUNE) \
                     .batch(4)

model = ...

model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), 
    loss="categorical_crossentropy", metrics=["acc"])

epochs = 200
model.fit(train_dataset, epochs=epochs, validation_data=validation_dataset)

 

위의 전체 코드에 대해서 순차적으로 설명드리겠습니다.

 

def read_image(mnist_dir=None):
  mnist_img_list = glob.glob('./mnist_png/**/*.png', recursive=True)
  mnist_img_array = np.array(mnist_img_list)
  label_array = np.array([os.path.dirname(os.path.basename(img_path)) for img_path in mnist_img_list])

  np.random.seed(200)
  np.random.shuffle(mnist_img_array)
  np.random.seed(200)
  np.random.shuffle(label_array)

  return mnist_img_array, to_categorical(label_array, 10)

 

(512, 512) 저장된 mnist_png 디렉토리 내 저장된 이미지들의 파일 경로들을 list -> np.array로 변환 및 shuffle 후 반환하는 코드입니다.

 

def fn(mnist_img_path, label):

  tmp = tf.io.read_file(mnist_img_path)
  decode_img = tf.image.decode_png(tmp, channels=1)
  decode_img = tf.cast(decode_img, tf.float32)

  return decode_img, tf.cast(label, tf.float32)

 

위의 함수는 mnist_img_path 및 label를 전달받아 Tensorflow 문자열 Tensor로 변환 후 tf.image.decode_png 함수를 통해 디코딩 과정을 거치게 됩니다. 최종적으로 tf.cast 함수를 통해 decode_img의 DType를 변환 후 반환합니다.

 

위의 코드에서 mnist_img_path(이미지)의 확장자가 .bmp, jpg인 경우에는 각각 다음 함수를 사용해주시면 됩니다.

 

• tf.image.decode_bmp()
• tf.image.decode_jpeg()

 

train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train)) \ 
                .map(fn, num_parallel_calls = tf.data.experimental.AUTOTUNE) \ 
                .shuffle(buffer_size=256) \ 
                .batch(4) \ 
                .prefetch(tf.data.experimental.AUTOTUNE)

 

위의 코드를 구성하는 요소들은 다음과 같습니다.

 

• map : 데이터셋의 각 요소에 주어진 fn 함수를 적용
• num_parallel_calls = tf.data.experimental.AUTOTUNE : 데이터 전처리를 병렬로 수행 (자동으로 최적의 Thread 수 결정)
• shuffle & buffer_size :  데이터셋을 무작위로 섞으며, 섞을 요소의 수를 정의
• prefetch & tf.data.experimental.AUT : 모델 학습 동안 다음 Batch를 미리 준비하여 입출력 시간 지연을 최소화하고, 계산과 데이터 로드를 병렬로 수행

---

이상으로 Tensorflow 2에서 대용량 데이터 학습하는 방법에 대해서 알아보았습니다.