pandas
tensorflow
dataframe
dataset conversion
machine learning

How to convert pandas dataframe to tensorflow dataset?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Converting a pandas DataFrame to a TensorFlow dataset is a common step when moving from exploratory data work into model training. The easiest path is usually tf.data.Dataset.from_tensor_slices, but the correct structure depends on whether you have labels, mixed feature types, or data too large to materialize comfortably in memory. The goal is not just conversion, but building a dataset that trains cleanly and predictably.

Convert Features and Labels Explicitly

For supervised learning, split the DataFrame into features and labels before conversion. TensorFlow works best when the dataset element shape matches the model input structure.

python
1import pandas as pd
2import tensorflow as tf
3
4df = pd.DataFrame({
5    "age": [21, 35, 52, 46],
6    "income": [40000, 72000, 98000, 65000],
7    "bought": [0, 1, 1, 0],
8})
9
10features = df[["age", "income"]]
11labels = df["bought"]
12
13dataset = tf.data.Dataset.from_tensor_slices((
14    dict(features),
15    labels
16))
17
18for x, y in dataset.take(2):
19    print(x, y)

Using dict(features) is convenient because each column becomes a named tensor, which maps naturally to Keras feature inputs.

Batch, Shuffle, and Prefetch

A raw dataset works, but most training pipelines also need batching and shuffling.

python
1train_ds = (
2    dataset
3    .shuffle(buffer_size=len(df), seed=42)
4    .batch(2)
5    .prefetch(tf.data.AUTOTUNE)
6)

This is where tf.data starts adding real value over passing arrays directly to model.fit.

Handle Numeric Arrays for Simpler Models

If your model expects a single dense tensor rather than named inputs, convert the feature frame to a NumPy array.

python
1import numpy as np
2
3x = features.to_numpy(dtype=np.float32)
4y = labels.to_numpy(dtype=np.float32)
5
6dense_ds = tf.data.Dataset.from_tensor_slices((x, y)).batch(2)
7
8for batch_x, batch_y in dense_ds.take(1):
9    print(batch_x.shape, batch_y.shape)

This is often the simplest path for dense tabular models.

Be Careful with String and Categorical Columns

Mixed dtypes can surprise you. String columns are supported, but many models need them encoded before training. You can either preprocess in pandas first or use Keras preprocessing layers.

python
1df = pd.DataFrame({
2    "city": ["Toronto", "Paris", "Toronto"],
3    "rooms": [2, 3, 1],
4    "label": [1, 0, 1],
5})
6
7feature_ds = tf.data.Dataset.from_tensor_slices((
8    {
9        "city": df["city"].astype(str).to_numpy(),
10        "rooms": df["rooms"].to_numpy(dtype="float32"),
11    },
12    df["label"].to_numpy(dtype="float32"),
13))

The conversion works, but the model still needs a strategy for turning "city" into useful numeric features.

Use a Generator Only When Necessary

from_tensor_slices is usually best for in-memory data. Use from_generator only when you truly need streaming behavior or custom row-by-row logic.

python
1def row_generator():
2    for _, row in df.iterrows():
3        yield {"rooms": float(row["rooms"])}, float(row["label"])
4
5gen_ds = tf.data.Dataset.from_generator(
6    row_generator,
7    output_signature=(
8        {"rooms": tf.TensorSpec(shape=(), dtype=tf.float32)},
9        tf.TensorSpec(shape=(), dtype=tf.float32),
10    ),
11)

This is more flexible, but also more verbose and slower if you use it without a real need.

Match Dataset Structure to Model Inputs

The most common source of errors is a mismatch between dataset elements and model input signatures. If the model expects named inputs, yield a dictionary. If it expects one tensor, yield one tensor. Keep that alignment explicit rather than relying on automatic coercion.

It is also worth validating one batch before training starts. A quick for loop over dataset.take(1) often reveals dtype or shape mistakes much faster than waiting for a model error inside fit.

Common Pitfalls

  • Passing the whole DataFrame directly without separating labels first.
  • Mixing object dtype columns into numeric pipelines without conversion.
  • Using from_generator when from_tensor_slices would be simpler and faster.
  • Forgetting to batch and prefetch before training.
  • Producing dataset elements whose structure does not match the model input signature.

Summary

  • Split DataFrame features and labels before conversion.
  • Use from_tensor_slices for the normal in-memory case.
  • Choose dictionary inputs or dense arrays based on model structure.
  • Clean up string and categorical columns deliberately before training.
  • Add batching, shuffling, and prefetching so the dataset is training-ready.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.