tf.data.Dataset Basics: Load Data From Tensors, CSV and Files

The Story — The Factory Conveyor Belt

So far we fed tensors straight into the model. Real projects need more: reading files, shuffling, batching, and doing it fast. This module covers the whole pipeline and the habits of good training. This first lesson is about where data comes from.

From Tensors and Arrays

from_tensor_slices cuts along the first axis: 5 rows become 5 elements. Tuples and dicts are cut in step, so each feature stays with its label.

import tensorflow as tf
import numpy as np

features = np.arange(15, dtype="float32").reshape(5, 3)   # 5 samples, 3 features
labels = np.array([0, 1, 0, 1, 1])

ds = tf.data.Dataset.from_tensor_slices((features, labels))
print(ds.element_spec)
print("number of elements:", ds.cardinality().numpy())
for x, y in ds.take(2):
    print("x =", x.numpy(), " y =", y.numpy())

dict_ds = tf.data.Dataset.from_tensor_slices({"age": [25, 32, 47], "city": ["Delhi", "Pune", "Baddi"]})
for row in dict_ds:
    print(int(row["age"]), row["city"].numpy().decode())

print("range      :", [int(v) for v in tf.data.Dataset.range(3, 9, 2)])

def gen():                                      # any Python generator
    for i in range(3):
        yield i, i * i
gen_ds = tf.data.Dataset.from_generator(gen, output_signature=(
    tf.TensorSpec([], tf.int32), tf.TensorSpec([], tf.int32)))
print("generator  :", [(int(a), int(b)) for a, b in gen_ds])

Try It — The Slicer

From CSV Files

First we write a small CSV file, so the example runs anywhere.

csv_text = """age,income,city,bought
25,32000,Delhi,0
32,58000,Pune,1
47,91000,Baddi,1
51,40000,Delhi,0
38,75000,Pune,1
29,28000,Baddi,0
"""
with open("/tmp/customers.csv", "w") as f:
    f.write(csv_text)

The Easy Way — make_csv_dataset

It reads the header, guesses each column's type, picks out the label, batches and shuffles.

csv_ds = tf.data.experimental.make_csv_dataset(
    "/tmp/customers.csv", batch_size=3, label_name="bought",
    num_epochs=1, shuffle=False)

for features, label in csv_ds.take(1):
    for name, col in features.items():
        print(f"{name:7s} {col.dtype.name:7s} {col.numpy()}")
    print("label  ", label.numpy())

The Manual Way — TextLineDataset + decode_csv

Useful when you want full control over column types and defaults for missing values.

defaults = [tf.constant(0.0), tf.constant(0.0), tf.constant(""), tf.constant(0)]   # one per column

def parse_line(line):
    age, income, city, bought = tf.io.decode_csv(line, record_defaults=defaults)
    x = tf.stack([age, income / 1000.0])          # numeric features
    return x, bought

lines = tf.data.TextLineDataset("/tmp/customers.csv").skip(1)   # skip the header
manual = lines.map(parse_line)
for x, y in manual.take(3):
    print(x.numpy(), y.numpy())

Try It — CSV Parser Playground

From Files and Folders

A very common layout stores one image per file, with the class as the folder name: data/cat/001.png. Let us make such a folder, then load it.

import os
for label in ["cat", "dog"]:
    os.makedirs(f"/tmp/pets/{label}", exist_ok=True)
    for i in range(3):
        img = tf.cast(tf.random.uniform([8, 8, 3], 0, 255, seed=i), tf.uint8)   # a tiny fake image
        tf.io.write_file(f"/tmp/pets/{label}/{i}.png", tf.io.encode_png(img))

class_names = tf.constant(["cat", "dog"])

def load(path):
    img = tf.io.decode_png(tf.io.read_file(path), channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)            # 0..1
    folder = tf.strings.split(path, os.sep)[-2]
    label = tf.argmax(tf.cast(folder == class_names, tf.int32))
    return img, label

files = tf.data.Dataset.list_files("/tmp/pets/*/*.png", shuffle=False)
images = files.map(load)
for img, lab in images.take(4):
    print(img.shape, img.dtype.name, "label", lab.numpy())
print("files found:", files.cardinality().numpy())

Which Source Should I Use?

Golden Rules