Datasets (tf.data)

import tensorflow as tf

# Canonical pipeline order:
# load → cache → shuffle → repeat → map → batch → prefetch
dataset = (
    tf.data.Dataset.from_tensor_slices((x, y))
    .cache()
    .shuffle(buffer_size=1000)
    .map(preprocess_fn, num_parallel_calls=tf.data.AUTOTUNE)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)

# From in-memory arrays (numpy or tensors)
ds = tf.data.Dataset.from_tensor_slices(tensor)           # one element per row
ds = tf.data.Dataset.from_tensor_slices((x, y))           # (features, labels) tuple
ds = tf.data.Dataset.from_tensor_slices({'img': x, 'label': y})  # dict

# From a generator
def gen():
    for i in range(100):
        yield i, i * 2

ds = tf.data.Dataset.from_generator(
    gen,
    output_signature=(
        tf.TensorSpec(shape=(), dtype=tf.int32),
        tf.TensorSpec(shape=(), dtype=tf.int32),
    )
)

# From tensors (wraps each as a single-element dataset)
ds = tf.data.Dataset.from_tensors((x, y))

# Range
ds = tf.data.Dataset.range(100)
ds = tf.data.Dataset.range(10, 100, 2)    # start, stop, step

# Read
raw = tf.data.TFRecordDataset('data.tfrecord')
raw = tf.data.TFRecordDataset(['file1.tfrecord', 'file2.tfrecord'])
raw = tf.data.TFRecordDataset(filenames, compression_type='GZIP')

# Parse
feature_desc = {
    'image':  tf.io.FixedLenFeature([], tf.string),
    'label':  tf.io.FixedLenFeature([], tf.int64),
    'values': tf.io.FixedLenSequenceFeature([], tf.float32, allow_missing=True),
}

def parse(example_proto):
    return tf.io.parse_single_example(example_proto, feature_desc)

ds = raw.map(parse, num_parallel_calls=tf.data.AUTOTUNE)

# Write TFRecords
def make_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])),
    }
    return tf.train.Example(features=tf.train.Features(feature=feature))

with tf.io.TFRecordWriter('out.tfrecord') as w:
    for img, lbl in zip(images, labels):
        w.write(make_example(img.tobytes(), lbl).SerializeToString())

# make_csv_dataset still lives under tf.data.experimental (no stable alias)
ds = tf.data.experimental.make_csv_dataset(
    'data.csv',
    batch_size=32,
    label_name='target',
    num_epochs=1,
    ignore_errors=True,
)

# Lower-level: one tensor tuple per CSV row
ds = tf.data.experimental.CsvDataset(
    'data.csv',
    record_defaults=[tf.float32, tf.float32, tf.int32],
    header=True,
)

ds = tf.data.TextLineDataset('data.txt')
ds = tf.data.TextLineDataset(['a.txt', 'b.txt']).skip(1)   # skip header

def load_image(path, label):
    img = tf.io.read_file(path)
    img = tf.image.decode_jpeg(img, channels=3)
    img = tf.image.resize(img, [224, 224])
    img = tf.cast(img, tf.float32) / 255.0
    return img, label

ds = tf.data.Dataset.from_tensor_slices((paths, labels)).map(load_image)

ds = ds.map(lambda x, y: (x / 255.0, y))
ds = ds.map(preprocess, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.map(preprocess, num_parallel_calls=4, deterministic=False)  # order not guaranteed

ds = ds.filter(lambda x, y: y != -1)
ds = ds.filter(lambda x: tf.math.reduce_sum(x) > 0)

# flat_map — map then flatten one level
ds = ds.flat_map(lambda x: tf.data.Dataset.from_tensor_slices(x))

# interleave — parallel read of nested datasets (use for files)
filenames = tf.data.Dataset.list_files('data/*.tfrecord', shuffle=True)
ds = filenames.interleave(
    lambda f: tf.data.TFRecordDataset(f),
    cycle_length=tf.data.AUTOTUNE,
    num_parallel_calls=tf.data.AUTOTUNE,
    deterministic=False,
)

ds.batch(32)
ds.batch(32, drop_remainder=True)    # drop last incomplete batch (needed for TPU)
ds.unbatch()                          # inverse of batch

# Pad variable-length sequences
ds.padded_batch(
    32,
    padded_shapes=([None], []),        # pad sequence dim, scalar label
    padding_values=(0, -1),
)

ds.shuffle(
    buffer_size=10000,    # larger = more random; 10x batch_size is common
    seed=42,
    reshuffle_each_iteration=True,   # re-shuffle each epoch when repeated
)

ds.repeat()          # infinite
ds.repeat(5)         # repeat 5 times

# Prefer passing epochs= to model.fit() over .repeat()
# but .repeat() is needed when steps_per_epoch is set manually

ds.cache()              # cache in memory (first epoch reads disk; subsequent use memory)
ds.cache('cache_dir/')  # cache to disk (persists across runs)

ds.prefetch(buffer_size=tf.data.AUTOTUNE)   # overlap preprocessing and model training
ds.prefetch(2)                               # manual buffer

ds.take(n)                  # first n elements
ds.skip(n)                  # skip first n
ds.enumerate()              # yields (index, element)
ds.zip((ds_a, ds_b))        # zip two datasets
tf.data.Dataset.zip((ds_a, ds_b))

ds.concatenate(ds2)         # append ds2 after ds
tf.data.Dataset.sample_from_datasets([ds1, ds2], weights=[0.7, 0.3])  # weighted mix
ds.window(size=5, shift=1)  # rolling windows
ds.reduce(0, lambda acc, x: acc + x)   # fold
ds.ignore_errors()          # skip bad records (stable since TF 2.9)
ds.ignore_errors(log_warning=True)     # log each skipped record

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
# Also: cifar10, cifar100, fashion_mnist, imdb, reuters, boston_housing

import tensorflow_datasets as tfds

ds, info = tfds.load('imagenet2012', split='train', with_info=True,
                     as_supervised=True)
ds = ds.map(lambda img, lbl: (tf.image.resize(img, [224, 224]), lbl))

# Splits
train_ds, val_ds = tfds.load('mnist', split=['train', 'test'], as_supervised=True)
train_80, train_20 = tfds.load('mnist', split=['train[:80%]', 'train[80%:]'])

# Full optimized pipeline for images
AUTOTUNE = tf.data.AUTOTUNE

train_ds = (
    tf.data.Dataset.list_files('train/*.tfrecord', shuffle=True)
    .interleave(
        lambda p: tf.data.TFRecordDataset(p),
        cycle_length=AUTOTUNE,
        num_parallel_calls=AUTOTUNE,
    )
    .map(parse_and_augment, num_parallel_calls=AUTOTUNE)
    .cache()
    .shuffle(1024)
    .batch(64, drop_remainder=True)
    .prefetch(AUTOTUNE)
)

ds.cardinality()          # number of elements (-1 if unknown, -2 if infinite)
ds.element_spec           # TensorSpec tree describing element structure

# Iterate
for batch in ds.take(3):
    print(batch)

# To numpy (small datasets only)
list(ds.as_numpy_iterator())