How to load CSV data into a TensorFlow dataset

CSV exports often sit between spreadsheet work, reporting systems, and TensorFlow training code. Loading a headered CSV file into a tf.data.Dataset keeps that handoff inside TensorFlow so feature columns, labels, batching, and later prefetching use one dataset contract.

tf.data.experimental.make_csv_dataset() reads one or more CSV files and returns batched dataset elements. When label_name is set, each element contains a feature dictionary keyed by column name plus a separate label tensor, which matches the input shape expected by supervised training code.

The local check uses num_epochs=1 and shuffle=False so the run ends after one pass through the file and prints the first batch in file order. Keep shuffling enabled for real training after the column mapping has been checked, and use tf.data.experimental.CsvDataset only when headerless files or fixed column defaults require lower-level control.

1. Save a headered CSV file as training.csv.

   training.csv

   feature_a,feature_b,label
   0.10,1.20,0
   0.40,0.70,1
   0.30,1.10,0
   0.90,0.20,1

   The header row becomes feature dictionary keys. The label column is removed from the feature dictionary when the loader uses label_name=“label”.
   Tool: Comma-Separated Values (CSV) Converter

2. Save the CSV loader as load_csv_dataset.py.

   load_csv_dataset.py

   from pprint import pprint
    
   import tensorflow as tf
    
   dataset = tf.data.experimental.make_csv_dataset(
       "training.csv",
       batch_size=2,
       label_name="label",
       num_epochs=1,
       shuffle=False,
   )
    
   first_features, first_labels = next(iter(dataset))
   row_count = sum(int(labels.shape[0]) for _, labels in dataset)
    
   rounded_features = {
       name: [round(float(value), 2) for value in tensor.numpy().tolist()]
       for name, tensor in first_features.items()
   }
    
   print(type(dataset).__name__)
   pprint(dataset.element_spec)
   print(rounded_features)
   print(first_labels.numpy().tolist())
   print(f"rows={row_count}")

   make_csv_dataset() infers column data types, batches rows, and wraps the result in a tf.data.Dataset pipeline. Set num_epochs=1 for finite verification runs so the dataset does not repeat indefinitely.

3. Run the loader and confirm that the feature keys, label tensor, first batch, and row count match the CSV file.

   $ python3 load_csv_dataset.py
   _PrefetchDataset
   (OrderedDict([('feature_a',
                  TensorSpec(shape=(None,), dtype=tf.float32, name=None)),
                 ('feature_b',
                  TensorSpec(shape=(None,), dtype=tf.float32, name=None))]),
    TensorSpec(shape=(None,), dtype=tf.int32, name=None))
   {'feature_a': [0.1, 0.4], 'feature_b': [1.2, 0.7]}
   [0, 1]
   rows=4

   TensorSpec(shape=(None,), dtype=tf.float32) shows that each feature column is batched as a numeric tensor. The final rows=4 line confirms that all four CSV records were read.

4. Remove the temporary CSV and loader files when they were created only for the local check.

   $ rm training.csv load_csv_dataset.py