TensorFlow
session management
machine learning
SavedModel
prediction optimization

How to keep tensorflow session open between predictions? Loading from SavedModel

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

If you load a SavedModel for every prediction, most of the latency comes from model loading rather than inference. In session-based TensorFlow code, the right pattern is to load the graph once, keep the session alive, and reuse it for repeated session.run calls. In modern TensorFlow, the equivalent idea is to load the model once and reuse the callable object, not recreate it per request.

Load the SavedModel Once in TensorFlow 1.x Style Code

When using TensorFlow session APIs, create the graph and session a single time during process startup, then keep references to the input and output tensors for future predictions.

python
1import numpy as np
2import tensorflow as tf
3
4tf.compat.v1.disable_eager_execution()
5
6class Predictor:
7    def __init__(self, export_dir: str):
8        self.graph = tf.Graph()
9        self.session = tf.compat.v1.Session(graph=self.graph)
10
11        with self.graph.as_default():
12            tf.compat.v1.saved_model.loader.load(
13                self.session,
14                [tf.saved_model.SERVING],
15                export_dir,
16            )
17
18            self.input_tensor = self.graph.get_tensor_by_name("serving_default_input:0")
19            self.output_tensor = self.graph.get_tensor_by_name("StatefulPartitionedCall:0")
20
21    def predict(self, values):
22        feed = {self.input_tensor: np.asarray(values, dtype=np.float32)}
23        return self.session.run(self.output_tensor, feed_dict=feed)
24
25    def close(self):
26        self.session.close()

This avoids rebuilding the graph and reopening the session on every request.

Reuse the Same Predictor Object

The pattern only works if the same predictor instance is reused across calls. For example:

python
1predictor = Predictor("/models/exported_model")
2
3print(predictor.predict([[1.0, 2.0, 3.0]]))
4print(predictor.predict([[4.0, 5.0, 6.0]]))
5
6predictor.close()

If the code constructs a new Predictor inside every request handler, the session is not really being kept open at all.

Keep Graph and Session Bound Together

In TensorFlow 1.x style code, graphs and sessions are tightly related. If you keep a session alive, keep the graph and tensor references that belong to it alive as well. Mixing tensor handles from one graph with a different session is a common source of runtime errors.

That means a clean wrapper object is usually better than scattering global variables such as sess, graph, and tensor names across a module.

Think About Concurrency

A long-lived session is efficient, but concurrency still matters. If multiple threads call the same predictor at once, the serving layer should be designed intentionally. In simple applications, one process-local predictor object is enough. In high-throughput systems, a dedicated serving stack, batching layer, or per-worker model instance may be more appropriate.

The main point is that reusing the session saves load time, but it does not automatically solve serving architecture.

Modern TensorFlow Uses Model Objects Instead of Sessions

In TensorFlow 2, you usually load the SavedModel once and keep the returned object in memory. The same performance principle applies, even though the API no longer revolves around Session.

python
1import tensorflow as tf
2
3model = tf.saved_model.load("/models/exported_model")
4infer = model.signatures["serving_default"]
5
6result_one = infer(input=tf.constant([[1.0, 2.0, 3.0]]))
7result_two = infer(input=tf.constant([[4.0, 5.0, 6.0]]))
8
9print(result_one)
10print(result_two)

The mistake to avoid is the same as before: do not call tf.saved_model.load for every prediction if the process can reuse the loaded model.

Close Resources Intentionally

A persistent session consumes memory, file handles, and device resources. That is fine for a long-running prediction service, but cleanup still matters for scripts, tests, or reloadable workers. Give the wrapper a close method or use application shutdown hooks so resources are released deliberately.

Common Pitfalls

  • Reloading the SavedModel on every prediction and paying startup cost repeatedly.
  • Reusing a session but not the graph and tensor references that belong to it.
  • Creating the predictor object inside each request handler instead of at process startup.
  • Assuming a long-lived session alone solves concurrency and serving throughput.
  • Applying session-based TensorFlow 1.x advice directly to TensorFlow 2 code without translating the pattern.

Summary

  • For session-based TensorFlow, load the SavedModel once and reuse the same session for repeated predictions.
  • Store the graph, session, and tensor handles together in one wrapper object.
  • Keep the predictor instance alive across requests if you want the session to stay open.
  • In TensorFlow 2, the equivalent pattern is to load the model once and reuse the loaded callable.
  • Clean up persistent model resources intentionally when the process shuts down.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the 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.