TensorFlow
estimator
model saving
machine learning
tutorial

How to save estimator in Tensorflow for later use?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

TensorFlow Estimators save training state automatically, but there are two different goals you might mean by "save for later use." One is continuing training or running prediction later from checkpoints in model_dir. The other is exporting a SavedModel for serving or deployment. Knowing which one you need prevents a lot of confusion.

Checkpoints in model_dir

An Estimator writes checkpoints to its model_dir while training. As long as you recreate the same Estimator configuration later and point it at the same directory, it can resume from the saved state.

python
1import tensorflow as tf
2
3feature_columns = [tf.feature_column.numeric_column("x", shape=(1,))]
4
5estimator = tf.estimator.DNNRegressor(
6    feature_columns=feature_columns,
7    hidden_units=[8, 8],
8    model_dir="./estimator_model",
9)
10
11def train_input_fn():
12    features = {"x": [[1.0], [2.0], [3.0], [4.0]]}
13    labels = [2.0, 4.0, 6.0, 8.0]
14    dataset = tf.data.Dataset.from_tensor_slices((features, labels))
15    return dataset.batch(2).repeat()
16
17estimator.train(input_fn=train_input_fn, steps=100)

After training, the ./estimator_model directory contains checkpoints and metadata. To use the model later, construct the same Estimator again with the same model_dir.

python
1import tensorflow as tf
2
3feature_columns = [tf.feature_column.numeric_column("x", shape=(1,))]
4
5restored_estimator = tf.estimator.DNNRegressor(
6    feature_columns=feature_columns,
7    hidden_units=[8, 8],
8    model_dir="./estimator_model",
9)
10
11def predict_input_fn():
12    features = {"x": [[5.0], [6.0]]}
13    dataset = tf.data.Dataset.from_tensor_slices(features)
14    return dataset.batch(2)
15
16for prediction in restored_estimator.predict(input_fn=predict_input_fn):
17    print(prediction["predictions"])

No separate manual restore call is needed; Estimator reads from model_dir automatically.

Exporting a SavedModel

If the goal is deployment, exporting a SavedModel is usually the better path. This creates a portable artifact for TensorFlow Serving or another inference environment.

python
1import tensorflow as tf
2
3feature_columns = [tf.feature_column.numeric_column("x", shape=(1,))]
4
5estimator = tf.estimator.DNNRegressor(
6    feature_columns=feature_columns,
7    hidden_units=[8, 8],
8    model_dir="./estimator_model",
9)
10
11def serving_input_receiver_fn():
12    feature_spec = {
13        "x": tf.compat.v1.placeholder(dtype=tf.float32, shape=[None, 1], name="x")
14    }
15    return tf.estimator.export.ServingInputReceiver(feature_spec, feature_spec)
16
17export_dir = estimator.export_saved_model(
18    export_dir_base="./exports",
19    serving_input_receiver_fn=serving_input_receiver_fn,
20)
21
22print(export_dir)

This creates a versioned export directory that can be loaded by serving systems later.

Which One Should You Use?

Use model_dir checkpoints when you want to:

  • continue training
  • evaluate later with the same Estimator code
  • run prediction from the same Python project

Use export_saved_model when you want to:

  • deploy the model to a serving environment
  • share a stable inference artifact
  • separate training code from production prediction code

The two approaches can coexist. Many projects keep checkpoints for training continuity and also export a SavedModel for deployment.

A Practical Caution About Estimator

Estimator is now considered a legacy TensorFlow API. It still appears in older codebases and can work well for maintenance tasks, but newer TensorFlow workflows usually prefer Keras and SavedModel-based export paths.

That means if you are starting a greenfield project, Estimator is probably not the first API to choose. If you are maintaining an existing Estimator pipeline, though, the checkpoint and export patterns above are still the right mental model.

Common Pitfalls

The most common mistake is expecting a special "save" method for checkpoints. Estimator already saves checkpoints in model_dir during training, so later reuse depends on rebuilding the Estimator with the same configuration.

Another mistake is changing the model definition while reusing an old model_dir. If feature columns or architecture change, the stored checkpoint may no longer match.

Developers also sometimes think checkpoint storage and serving export are the same thing. They are not. Checkpoints are training-state artifacts, while SavedModel exports are inference-oriented artifacts.

Finally, be careful with input signatures when exporting. A serving export that does not match the prediction request format is difficult to use later even if the checkpoint itself is valid.

Summary

  • Estimators save checkpoints automatically inside model_dir.
  • Recreate the same Estimator with the same model_dir to reuse it later.
  • Use export_saved_model when you need a deployment artifact.
  • Checkpoints and SavedModel exports serve different purposes.
  • Estimator is legacy today, but these patterns still matter in older TensorFlow codebases.

Course illustration
Course illustration

All Rights Reserved.