TFRecord and Interleave — Large-Scale Pipelines

What ImageNet-scale datasets actually look like

For datasets too large to fit in memory and stored across many files, TFRecord is TensorFlow's native binary format. Each file holds serialized tf.train.Example protocol buffers — read sequentially fast, decoded in parallel.

interleave reads multiple TFRecord files concurrently, mixing their elements. Combined with cycle_length=AUTOTUNE and num_parallel_calls=AUTOTUNE, this saturates I/O on slow storage (HDDs, network drives) where reading any single file would be the bottleneck.

When to use TFRecord: dataset size > ~5GB, file count > 100, or storage on a network filesystem. For CIFAR-10 / MNIST sizes, TFRecord is overkill.

Code

TFRecord 쓰기·python

import tensorflow as tf

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

with tf.io.TFRecordWriter("train_000.tfrecord") as writer:
    for image, label in zip(images, labels):
        writer.write(serialize_example(image, label))

Interleave로 읽기 — 병렬 I/O·python

import tensorflow as tf

AUTOTUNE = tf.data.AUTOTUNE

def parse_tfrecord(serialized):
    desc = {
        'image': tf.io.FixedLenFeature([], tf.string),
        'label': tf.io.FixedLenFeature([], tf.int64),
    }
    ex = tf.io.parse_single_example(serialized, desc)
    image = tf.io.decode_jpeg(ex['image'], channels=3)
    image = tf.cast(image, tf.float32) / 255.0
    image = tf.image.resize(image, [224, 224])
    return image, ex['label']

filenames = tf.data.Dataset.list_files("train_*.tfrecord")

ds = (
    filenames
    .interleave(
        lambda fn: tf.data.TFRecordDataset(fn),
        cycle_length=AUTOTUNE,
        num_parallel_calls=AUTOTUNE,
        deterministic=False,    # slight perf gain if order doesn't matter
    )
    .map(parse_tfrecord, num_parallel_calls=AUTOTUNE)
    .cache()
    .shuffle(buffer_size=10000)
    .batch(128, drop_remainder=True)
    .prefetch(AUTOTUNE)
)

TFRecord and Interleave — Large-Scale Pipelines

What ImageNet-scale datasets actually look like

Code

Progress

Comments 0