Introduction
Models trained with tf.estimator are saved automatically to the model_dir as checkpoints, or explicitly exported as SavedModel format using estimator.export_saved_model(). To reload and predict, either recreate the estimator with the same model_dir and call estimator.predict(), or load the exported SavedModel with tf.saved_model.load() for framework-independent inference. The SavedModel approach is preferred for deployment because it does not require the original model code. Note that tf.estimator is deprecated in TensorFlow 2.x — Keras model.save() and tf.saved_model are the recommended alternatives for new projects.
Training and Saving with tf.estimator
1import tensorflow as tf
2import numpy as np
3
4# Define feature columns
5feature_columns = [
6 tf.feature_column.numeric_column("x", shape=[4])
7]
8
9# Create the estimator
10estimator = tf.estimator.DNNClassifier(
11 feature_columns=feature_columns,
12 hidden_units=[128, 64],
13 n_classes=3,
14 model_dir="./my_model" # Checkpoints saved here automatically
15)
16
17# Training input function
18def train_input_fn():
19 dataset = tf.data.Dataset.from_tensor_slices((
20 {"x": np.random.randn(100, 4).astype(np.float32)},
21 np.random.randint(0, 3, 100)
22 ))
23 return dataset.batch(32).repeat(10)
24
25# Train — checkpoints are saved to model_dir during training
26estimator.train(input_fn=train_input_fn, steps=200)
Checkpoints in model_dir contain the model weights but require the original estimator code to reload. For portable deployment, export as SavedModel.
Exporting as SavedModel
1# Define a serving input function
2def serving_input_receiver_fn():
3 feature_spec = {
4 "x": tf.io.FixedLenFeature([4], tf.float32)
5 }
6 # For JSON/REST API inputs
7 serialized_tf_example = tf.compat.v1.placeholder(
8 dtype=tf.string, shape=[None], name="input_example_tensor"
9 )
10 receiver_tensors = {"examples": serialized_tf_example}
11 features = tf.io.parse_example(serialized_tf_example, feature_spec)
12 return tf.estimator.export.ServingInputReceiver(features, receiver_tensors)
13
14# Export SavedModel
15export_dir = estimator.export_saved_model(
16 export_dir_base="./exported_model",
17 serving_input_receiver_fn=serving_input_receiver_fn
18)
19print(f"Model exported to: {export_dir}")
20# ./exported_model/1234567890/ (timestamped subdirectory)
1# Simpler: raw feature input (no tf.Example parsing)
2def raw_serving_input_fn():
3 feature_placeholders = {
4 "x": tf.compat.v1.placeholder(tf.float32, [None, 4])
5 }
6 return tf.estimator.export.ServingInputReceiver(
7 feature_placeholders, feature_placeholders
8 )
9
10export_dir = estimator.export_saved_model(
11 export_dir_base="./exported_model_raw",
12 serving_input_receiver_fn=raw_serving_input_fn
13)
Method 1: Predict with the Same Estimator
1# Recreate estimator pointing to the same model_dir
2estimator = tf.estimator.DNNClassifier(
3 feature_columns=feature_columns,
4 hidden_units=[128, 64],
5 n_classes=3,
6 model_dir="./my_model" # Loads latest checkpoint automatically
7)
8
9# Prediction input function
10def predict_input_fn():
11 new_data = {
12 "x": np.array([
13 [5.1, 3.5, 1.4, 0.2],
14 [6.7, 3.1, 4.7, 1.5],
15 [7.2, 3.0, 5.8, 1.8],
16 ], dtype=np.float32)
17 }
18 dataset = tf.data.Dataset.from_tensor_slices(new_data)
19 return dataset.batch(3)
20
21# Run predictions
22predictions = estimator.predict(input_fn=predict_input_fn)
23
24for pred in predictions:
25 class_id = pred["class_ids"][0]
26 probability = pred["probabilities"][class_id]
27 print(f"Predicted class: {class_id}, probability: {probability:.4f}")
This method requires the same estimator code (feature columns, hidden units, etc.) and access to the model_dir. It is suitable for development and testing.
Method 2: Load SavedModel with tf.saved_model
1# Load the exported SavedModel
2loaded_model = tf.saved_model.load(str(export_dir))
3
4# Get the prediction function
5infer = loaded_model.signatures["serving_default"]
6
7# Prepare input
8input_data = tf.constant([
9 [5.1, 3.5, 1.4, 0.2],
10 [6.7, 3.1, 4.7, 1.5],
11], dtype=tf.float32)
12
13# Run prediction
14result = infer(x=input_data)
15print(result.keys()) # dict_keys(['class_ids', 'probabilities', ...])
16print(result["class_ids"].numpy())
17print(result["probabilities"].numpy())
1# If the serving function expects tf.Example format
2import tensorflow as tf
3
4def make_example(features):
5 feature_dict = {
6 "x": tf.train.Feature(
7 float_list=tf.train.FloatList(value=features)
8 )
9 }
10 example = tf.train.Example(
11 features=tf.train.Features(feature=feature_dict)
12 )
13 return example.SerializeToString()
14
15examples = [make_example([5.1, 3.5, 1.4, 0.2])]
16result = infer(examples=tf.constant(examples))
Method 3: Using TensorFlow Serving
1# Start TensorFlow Serving with the exported model
2docker run -p 8501:8501 \
3 --mount type=bind,source=$(pwd)/exported_model,target=/models/my_model \
4 -e MODEL_NAME=my_model \
5 tensorflow/serving
6
7# Make predictions via REST API
8curl -d '{"instances": [{"x": [5.1, 3.5, 1.4, 0.2]}]}' \
9 -X POST http://localhost:8501/v1/models/my_model:predict
1# Python client for TensorFlow Serving
2import requests
3import json
4
5data = {
6 "instances": [
7 {"x": [5.1, 3.5, 1.4, 0.2]},
8 {"x": [6.7, 3.1, 4.7, 1.5]},
9 ]
10}
11
12response = requests.post(
13 "http://localhost:8501/v1/models/my_model:predict",
14 json=data
15)
16predictions = response.json()["predictions"]
17for pred in predictions:
18 print(f"Class: {pred['class_ids'][0]}, Prob: {max(pred['probabilities']):.4f}")
Migrating to Keras (Recommended)
1# tf.estimator is deprecated — use Keras for new projects
2model = tf.keras.Sequential([
3 tf.keras.layers.Dense(128, activation='relu', input_shape=(4,)),
4 tf.keras.layers.Dense(64, activation='relu'),
5 tf.keras.layers.Dense(3, activation='softmax')
6])
7
8model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
9model.fit(x_train, y_train, epochs=10)
10
11# Save
12model.save("./keras_model")
13
14# Load and predict — no model code needed
15loaded = tf.keras.models.load_model("./keras_model")
16predictions = loaded.predict(np.array([[5.1, 3.5, 1.4, 0.2]]))
17print(predictions)
Common Pitfalls
Changing model architecture and pointing to old model_dir: If you modify hidden_units, feature_columns, or n_classes and then point to an existing model_dir, the estimator fails to restore because the checkpoint shape does not match the new architecture. Delete the old model_dir or use a new path.
Forgetting the serving input function signature: export_saved_model requires a serving_input_receiver_fn that defines how the model receives input. Omitting it or mismatching the feature names causes export to fail or inference to receive wrong data.
Not finding the correct export subdirectory: export_saved_model creates a timestamped subdirectory (e.g., exported_model/1709234567/). Load the specific subdirectory, not the parent: tf.saved_model.load("exported_model/1709234567").
Using estimator.predict() without an input function: Unlike Keras model.predict(data), estimator requires an input function that returns a tf.data.Dataset. Passing raw arrays directly raises an error. Wrap data in tf.data.Dataset.from_tensor_slices().
Ignoring the tf.estimator deprecation: tf.estimator is deprecated since TensorFlow 2.0 and will be removed in a future version. For new models, use tf.keras with model.save() and tf.keras.models.load_model(), which provide simpler saving, loading, and prediction workflows.
Summary
Recreate the estimator with the same model_dir and call estimator.predict() for quick testing
Export with estimator.export_saved_model() and load with tf.saved_model.load() for deployment
Use loaded_model.signatures["serving_default"] to access the prediction function from a SavedModel
Deploy via TensorFlow Serving for production REST/gRPC inference
Migrate to tf.keras for new projects — model.save() and load_model() are simpler and actively maintained
Always use the same feature names and shapes in the serving input function as were used during training
