Introduction

If you are still maintaining TensorFlow Estimator code, batched training is done through an input_fn that returns a tf.data.Dataset. The old fit() style APIs are long gone from mainstream TensorFlow examples, and Estimator itself is now legacy: TensorFlow 2.15 was the final main release that included tf-estimator, and new code should generally prefer Keras.

What Estimator Expects

An Estimator does not read NumPy arrays directly through a fit(batch_size=...) style interface the way Keras does. Instead, train() calls an input_fn, and that function returns a dataset that yields batches of features and labels.

At a high level, the batch logic lives here:

• build a dataset
• shuffle if training
• batch it
• repeat if needed

Minimal input_fn Example

Here is a small example using NumPy arrays and a dataset pipeline.

1import tensorflow as tf
2import numpy as np
3
4X_train = np.array([[1.0], [2.0], [3.0], [4.0]], dtype=np.float32)
5y_train = np.array([2.0, 4.0, 6.0, 8.0], dtype=np.float32)
6
7def train_input_fn():
8    dataset = tf.data.Dataset.from_tensor_slices(({"x": X_train}, y_train))
9    dataset = dataset.shuffle(4).batch(2).repeat()
10    return dataset

This dataset yields batches of size 2 forever because of repeat(). Estimator training uses the steps argument to decide when to stop.

A Simple Estimator

The following example uses a premade linear estimator so the data-feeding pattern stays clear.

1import tensorflow as tf
2
3feature_columns = [tf.feature_column.numeric_column("x", shape=(1,))]
4estimator = tf.estimator.LinearRegressor(feature_columns=feature_columns)
5
6estimator.train(input_fn=train_input_fn, steps=100)

The crucial point is that batch size belongs to the dataset pipeline, not to a fit() argument.

Evaluation Input Functions Usually Differ

Training input pipelines often shuffle and repeat. Evaluation pipelines usually should not.

1import tensorflow as tf
2import numpy as np
3
4X_eval = np.array([[5.0], [6.0]], dtype=np.float32)
5y_eval = np.array([10.0, 12.0], dtype=np.float32)
6
7def eval_input_fn():
8    dataset = tf.data.Dataset.from_tensor_slices(({"x": X_eval}, y_eval))
9    dataset = dataset.batch(2)
10    return dataset

That difference matters because shuffling and infinite repetition are usually training-only behaviors.

Custom model_fn Still Uses the Same Input Pattern

If you are using a custom Estimator, the input_fn story does not change.

1import tensorflow as tf
2
3def model_fn(features, labels, mode):
4    x = features["x"]
5    predictions = tf.keras.layers.Dense(1)(x)
6
7    if mode == tf.estimator.ModeKeys.PREDICT:
8        return tf.estimator.EstimatorSpec(mode, predictions=predictions)
9
10    loss = tf.reduce_mean(tf.square(predictions[:, 0] - labels))
11
12    if mode == tf.estimator.ModeKeys.TRAIN:
13        optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.01)
14        train_op = optimizer.minimize(
15            loss, tf.compat.v1.train.get_or_create_global_step()
16        )
17        return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
18
19    return tf.estimator.EstimatorSpec(mode, loss=loss)

Whether the model is premade or custom, the batch feeding still comes from the dataset returned by input_fn.

Estimator Is Legacy Now

This part matters in 2026. Estimator is maintained for backwards compatibility, not as the primary TensorFlow training path. If you are starting fresh, Keras with model.fit(...) and tf.data.Dataset is usually the better choice.

Still, many production systems still contain Estimator code, so understanding the input_fn pattern is valuable for maintenance and migration.

Common Pitfalls

The biggest pitfall is looking for a fit(batch_size=...) style argument on Estimator. That is a Keras mental model, not an Estimator one.

Another issue is forgetting repeat() in the training dataset while also passing many training steps. If the dataset runs out early, training stops sooner than expected.

Developers also forget to remove shuffle() and repeat() from evaluation or prediction pipelines, which can make results harder to interpret.

Finally, be careful about TensorFlow version assumptions. Estimator examples from older blog posts often rely on APIs that have moved, been deprecated, or require the separate tf-estimator compatibility path.

Summary

• Estimator training uses train(input_fn=...), not a Keras-style fit(batch_size=...) interface.
• Put batching, shuffling, and repetition inside a tf.data.Dataset returned by input_fn.
• Use different dataset behavior for training and evaluation.
• The same input_fn pattern applies to premade and custom Estimators.
• Estimator is a legacy TensorFlow API, so new projects should usually prefer Keras.