본문 바로가기
Artificial Intelligence/60. Python

[PYTHON] 대규모 이미지 처리를 위한 TFRecord와 WebDataset 활용 방법 7가지 및 성능 차이 분석

by Papa Martino V 2026. 4. 19.
728x90

TFRecord와 WebDataset 활용
TFRecord와 WebDataset 활용

 

딥러닝 모델의 규모가 커짐에 따라 수백만 장 이상의 이미지를 학습에 활용하는 것이 일반화되었습니다. 하지만 수많은 작은 이미지 파일을 개별적으로 읽어오는 방식은 운영체제의 I/O 병목 현상을 야기하며, GPU가 연산보다 데이터를 기다리는 시간이 더 길어지는 비효율을 초래합니다. 이를 해결하기 위해 데이터를 거대한 바이너리 컨테이너로 묶는 전략이 필수적입니다.

본 포스팅에서는 TensorFlow 생태계의 표준인 TFRecord와 PyTorch 및 클라우드 네이티브 환경에서 각광받는 WebDataset(POSIX tar 기반)의 구조적 차이를 심층 비교하고, 실무 개발자가 즉시 적용 가능한 7가지 최적화 예시를 공유합니다.


1. 데이터 직렬화 포맷의 핵심 차이점과 선택 해결 기준

데이터셋을 구성할 때 어떤 포맷을 선택하느냐는 인프라 환경(로컬 서버 vs 클라우드 버킷)과 사용하는 프레임워크에 따라 달라집니다.

비교 항목 TFRecord (TensorFlow) WebDataset (PyTorch/General) 성능 해결 시사점
기본 구조 Protobuf 기반 바이너리 POSIX 표준 .tar 파일 표준성 및 호환성 차이
접근 방식 시퀀셜 읽기 중심 샤드(Shard) 단위 스트리밍 클라우드 스토리지 최적화
데이터 디코딩 엄격한 스키마 정의 필요 파일 확장자 기반 자동 매핑 유연한 데이터 구조 변경
셔플링 성능 버퍼 기반 셔플링 샤드 혼합 및 버퍼링 결합 대규모 데이터 무작위성 확보
의존성 TensorFlow 필수 표준 라이브러리 및 가벼운 패키지 프레임워크 독립적 설계 가능

2. 실무를 위한 고성능 데이터 파이프라인 구현 패턴 (7 Examples)

현업에서 직면하는 데이터 로딩 병목을 해결하기 위한 파이썬 코드 예시입니다.

Example 1: 이미지와 메타데이터를 TFRecord로 직렬화하기

이미지 바이너리와 정수형 라벨을 하나의 Protocol Buffer 메시지로 묶는 가장 기본적인 방법입니다.

import tensorflow as tf

def serialize_example(image_bytes, label):
    feature = {
        'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_bytes])),
        'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
    }
    example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
    return example_proto.SerializeToString()

with tf.io.TFRecordWriter("dataset.tfrecord") as writer:
    img_data = open("sample.jpg", "rb").read()
    example = serialize_example(img_data, 1)
    writer.write(example)

Example 2: tf.data API를 활용한 TFRecord 고속 파싱 및 프리패치

CPU에서 압축을 풀고 GPU로 데이터를 미리 밀어 넣어 GPU 유휴 시간을 제거하는 해결 방법입니다.

def parse_func(example_proto):
    feature_description = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([], tf.int64),
    }
    parsed = tf.io.parse_single_example(example_proto, feature_description)
    image = tf.io.decode_jpeg(parsed['image'], channels=3)
    return image, parsed['label']

dataset = tf.data.TFRecordDataset("dataset.tfrecord")
dataset = dataset.map(parse_func, num_parallel_calls=tf.data.AUTOTUNE)
dataset = dataset.batch(32).prefetch(buffer_size=tf.data.AUTOTUNE)

Example 3: WebDataset 형식으로 대규모 샤드 생성하기

수백 기가바이트의 데이터를 1GB 단위의 .tar 파일(샤드)로 분할하여 저장하는 방식입니다.

import webdataset as wds

writer = wds.ShardWriter("shards-%06d.tar", maxsize=1e9) # 1GB마다 새 파일 생성
for i in range(1000):
    with open(f"img_{i}.jpg", "rb") as f:
        image_data = f.read()
    writer.write({
        "__key__": f"sample_{i}",
        "jpg": image_data,
        "cls": i % 10
    })
writer.close()

Example 4: PyTorch에서 WebDataset 스트리밍 로더 구현

클라우드 저장소(S3 등)에서 로컬 디스크 복사 없이 바로 데이터를 스트리밍하여 학습하는 방법입니다.

dataset = (
    wds.WebDataset("shards-{000000..000010}.tar")
    .shuffle(1000)
    .decode("pil")
    .to_tuple("jpg", "cls")
)
loader = torch.utils.data.DataLoader(dataset, batch_size=64, num_workers=4)

Example 5: TFRecord 다중 파일 병렬 읽기로 처리량(Throughput) 해결

하나의 거대한 파일 대신 여러 파일에서 동시에 읽어 병목을 해소합니다.

files = tf.data.Dataset.list_files("train-*.tfrecord")
dataset = files.interleave(
    lambda x: tf.data.TFRecordDataset(x),
    cycle_length=8,
    num_parallel_calls=tf.data.AUTOTUNE
)

Example 6: WebDataset을 이용한 데이터 증강(Augmentation) 파이프라인 결합

로딩 과정에 전처리 로직을 심리스하게 결합하는 예시입니다.

def preprocess(sample):
    image, label = sample
    # torchvision transform 적용
    return transform(image), label

dataset = dataset.map(preprocess)

Example 7: 가변 길이 데이터(예: 비디오 프레임) 처리를 위한 시퀀스 저장 해결

TFRecord의 SequenceExample을 활용하여 가변적 데이터를 구조화하는 방법입니다.

def make_sequence_example(frames, label):
    frame_list = [tf.train.Feature(bytes_list=tf.train.BytesList(value=[f])) for f in frames]
    feature_list = tf.train.FeatureList(feature=frame_list)
    
    return tf.train.SequenceExample(
        context=tf.train.Features(feature={'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))}),
        feature_lists=tf.train.FeatureLists(feature_list={'frames': feature_list})
    ).SerializeToString()

3. 대규모 데이터 처리 성능 최적화를 위한 3대 원칙

단순히 포맷을 바꾼다고 성능이 올라가지는 않습니다. 다음 3가지 원칙을 지켜야 합니다.

  • 샤딩(Sharding): 단일 파일의 크기를 100MB~1GB 정도로 유지하여 네트워크 타임아웃을 방지하고 병렬성을 확보하십시오.
  • 로컬 캐싱: WebDataset 사용 시 원격 저장소의 데이터를 로컬 노드 메모리나 고속 SSD에 임시 캐싱하여 반복 학습 속도를 높이십시오.
  • 비동기 전처리: num_workersprefetch 설정을 통해 다음 배치를 GPU가 연산하는 동안 미리 CPU에서 준비하도록 설계하십시오.

4. 결론 및 향후 전망

TensorFlow 환경에 고착되어 있다면 TFRecord가 최선이지만, 유연성과 멀티 프레임워크 호환성, 클라우드 확정성을 고려한다면 WebDataset이 강력한 대안입니다. 데이터 파이프라인의 설계 능력이 곧 모델의 학습 속도와 직결되는 만큼, 인프라의 물리적 특성을 이해하고 최적의 포맷을 선택하는 안목이 필요합니다.

 

내용 출처 및 참조:

  • TensorFlow Official Guide: "TFRecord and tf.train.Example" (2025)
  • WebDataset Documentation: "Efficient storage and loading of deep learning datasets"
  • High-Performance Data Loading in PyTorch (NVIDIA Developer Blog)
728x90