How to create a dataset from tensors in TensorFlow

Creating a tf.data.Dataset from tensors is the quickest way to turn small in-memory feature arrays, label arrays, or dictionaries into an input pipeline that TensorFlow can iterate. This fits workflows where the data is already loaded in Python and each row should become one dataset element.

tf.data.Dataset.from_tensor_slices() slices every input tensor along its first dimension and preserves the input structure. A 2D feature tensor becomes a dataset of row tensors, while a tuple or dictionary of tensors becomes dataset elements with matching feature and label components.

All input components must have the same size in their first dimension because TensorFlow uses that leading axis as the dataset length. This pattern is best for in-memory data; when the source is too large to keep as tensors or should stream from files, move to a file-backed tf.data pipeline instead of building the dataset from tensor slices.

1. Open a terminal in a Python environment where TensorFlow already imports cleanly.

   $ python3 - <<'PY'
   import tensorflow as tf
   print(tf.__version__)
   PY
   2.21.0

   Related: How to create a virtual environment for TensorFlow
   Related: How to install TensorFlow with pip

2. Save the demo as

   tensor_slices_demo.py

   so the feature dictionary and label tensor share the same row index.

   tensor_slices_demo.py

   import tensorflow as tf
    
   features = {
       "sepal_length": tf.constant([5.1, 4.9, 6.7, 5.6], dtype=tf.float32),
       "sepal_width": tf.constant([3.5, 3.0, 3.1, 2.8], dtype=tf.float32),
   }
   labels = tf.constant([0, 0, 2, 1], dtype=tf.int32)
    
   dataset = tf.data.Dataset.from_tensor_slices((features, labels))
   batched = dataset.batch(2)
    
   first_features, first_label = next(iter(dataset))
   batch_features, batch_labels = next(iter(batched))
   feature_specs, label_spec = dataset.element_spec
    
   rounded_first_features = {
       name: round(float(value.numpy()), 1) for name, value in first_features.items()
   }
   rounded_batch_features = {
       name: [round(float(item), 1) for item in value.numpy().tolist()]
       for name, value in batch_features.items()
   }
   print(f"tensorflow={tf.__version__}")
   print(f"dataset_cardinality={int(dataset.cardinality())}")
   print(
       "sepal_length_spec="
       f"shape={feature_specs['sepal_length'].shape}, "
       f"dtype={feature_specs['sepal_length'].dtype.name}"
   )
   print(
       "sepal_width_spec="
       f"shape={feature_specs['sepal_width'].shape}, "
       f"dtype={feature_specs['sepal_width'].dtype.name}"
   )
   print(f"label_spec=shape={label_spec.shape}, dtype={label_spec.dtype.name}")
   print(f"first_example={rounded_first_features}, label={first_label.numpy().item()}")
   print(f"first_batch={rounded_batch_features}, labels={batch_labels.numpy().tolist()}")

   The feature input is a Python dictionary of tensors, so each dataset element keeps the same dictionary keys instead of flattening them into one unnamed tensor.

3. Build the dataset with Dataset.from_tensor_slices¹⁾ so TensorFlow slices the first dimension of both inputs into aligned examples.

   dataset = tf.data.Dataset.from_tensor_slices((features, labels))

   Every tensor or dictionary value must have the same number of rows in axis 0 or TensorFlow raises a shape mismatch error when the dataset is created.

4. Run the script and confirm the dataset reports scalar per-example specs before the first batch groups two rows together.

   $ python3 tensor_slices_demo.py
   tensorflow=2.21.0
   dataset_cardinality=4
   sepal_length_spec=shape=(), dtype=float32
   sepal_width_spec=shape=(), dtype=float32
   label_spec=shape=(), dtype=int32
   first_example={'sepal_length': 5.1, 'sepal_width': 3.5}, label=0
   first_batch={'sepal_length': [5.1, 4.9], 'sepal_width': [3.5, 3.0]}, labels=[0, 0]

5. Extend the same dataset into the training input pipeline by adding shuffle, batching, and prefetch after the slice step.

   train_ds = (tf.data.Dataset.from_tensor_slices((features, labels))
       .shuffle(len(labels))
       .batch(32)
       .prefetch(tf.data.AUTOTUNE))

   Related: How to optimize TensorFlow data pipeline performance
   Related: How to train, evaluate, and run prediction with a Keras model

6. Remove the demo script when it was created only to check the tensor-slice layout.

   $ rm tensor_slices_demo.py