TensorFlow 대용량 데이터 처리

100GB가 넘는 이미지 데이터를 처리해야 한다. 이 정도 데이터라면 절대 한 번에 읽고 처리할 수 없다. 따라서 데이터를 샘플 단위로 나눠서 I/O 작업 효율을 높여야 한다. Tensorflow에서 사용하는 아이디어는 다음과 같다.

• 바이트를 순차적(sequantial)으로 접근할 수 있도록 저장한다.
• 한 번에 읽을 수 있는 단위로 나누어 저장한다.
• I/O 작업을 병렬로 처리해 효율성을 높인다.

데이터 저장

import tensorflow as tf

# 🔥이미지를 byte로 직렬화하고 수치 정보를 포함하여 하나의 Example로 만든다.
def serialize_example(image, numeric_features):
    feature = {
        'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image.tobytes()])),
        'features': tf.train.Feature(float_list=tf.train.FloatList(value=numeric_features)),
    }

    # 하나의 Example 객체 생성
    example = tf.train.Example(features=tf.train.Features(feature=feature))
    
    # 직렬화하여 TFRecord에 쓸 수 있는 binary 형식으로 변환
    return example.SerializeToString()


# 🔥TFRecordWriter를 사용하여 파일로 저장
with tf.io.TFRecordWriter('data_000.tfrecord') as writer:
    for image, numeric_features in dataset:
        serialized = serialize_example(image, numeric_features)
        writer.write(serialized)

데이터 불러오기

# TFRecord 안의 Example을 파싱하는 함수
def parse_example(example_proto):
    features = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'features': tf.io.FixedLenFeature([NUM_FEATURES], tf.float32)
    }

    # 🔥binary example을 dictionary 형태로 디코딩
    parsed = tf.io.parse_single_example(example_proto, features)
    image = tf.io.decode_raw(parsed['image'], tf.uint8)
    return image, parsed['features']


# 🔥TFRecord 파일 목록 불러오기 (여러 개 가능)
filenames = tf.io.gfile.glob("data_*.tfrecord")

# 🔥여러 TFRecord 파일로부터 데이터를 스트리밍 로드
dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=tf.data.AUTOTUNE)

# 🔥파싱 함수 적용 (병렬로 실행)
dataset = dataset.map(parse_example, num_parallel_calls=tf.data.AUTOTUNE)

# 배치 단위로 묶기
dataset = dataset.batch(BATCH_SIZE)

# 🔥I/O와 학습을 병렬화하여 처리 속도 개선
dataset = dataset.prefetch(tf.data.AUTOTUNE)

prefetch는 GPU가 학습 중일 때 다음 배치를 미리 준비해 병목 현상을 방지한다.