Practice Now
TensorFlow
Python
Machine Learning
Programming
Neural Networks

Tensorflow Creating a graph in a class and running it outside

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
TensorFlow & KerasMachine LearningPython

Introduction

In TensorFlow 1.x, creating a computation graph inside a class and running it outside required careful management of graphs and sessions. TensorFlow 2.x eliminates most of this complexity with eager execution and tf.function. This article covers both the legacy TF1 approach and the modern TF2 approach for encapsulating models in classes.

TensorFlow 1.x Approach (Legacy)

In TF1, you had to explicitly manage graphs and sessions:

python
1import tensorflow.compat.v1 as tf
2tf.disable_eager_execution()
3
4class MyModel:
5    def __init__(self, input_size, output_size):
6        self.graph = tf.Graph()
7
8        with self.graph.as_default():
9            self.x = tf.placeholder(tf.float32, shape=[None, input_size], name='input')
10            self.W = tf.Variable(tf.random.normal([input_size, output_size]), name='weights')
11            self.b = tf.Variable(tf.zeros([output_size]), name='bias')
12
13            self.output = tf.matmul(self.x, self.W) + self.b
14            self.prediction = tf.nn.softmax(self.output, name='prediction')
15
16            self.init_op = tf.global_variables_initializer()
17
18    def predict(self, data):
19        with tf.Session(graph=self.graph) as sess:
20            sess.run(self.init_op)
21            return sess.run(self.prediction, feed_dict={self.x: data})
22
23# Usage
24model = MyModel(input_size=10, output_size=3)
25import numpy as np
26result = model.predict(np.random.randn(5, 10))
27print(result.shape)  # (5, 3)

Persistent Session

For repeated calls, keep the session alive:

python
1class MyModelWithSession:
2    def __init__(self, input_size, output_size):
3        self.graph = tf.Graph()
4
5        with self.graph.as_default():
6            self.x = tf.placeholder(tf.float32, [None, input_size])
7            self.W = tf.Variable(tf.random.normal([input_size, output_size]))
8            self.b = tf.Variable(tf.zeros([output_size]))
9            self.output = tf.nn.softmax(tf.matmul(self.x, self.W) + self.b)
10            self.init_op = tf.global_variables_initializer()
11
12        self.session = tf.Session(graph=self.graph)
13        self.session.run(self.init_op)
14
15    def predict(self, data):
16        return self.session.run(self.output, feed_dict={self.x: data})
17
18    def close(self):
19        self.session.close()
20
21# Usage
22model = MyModelWithSession(10, 3)
23result = model.predict(np.random.randn(5, 10))
24model.close()

TF2 uses eager execution by default. No graphs or sessions to manage:

Using tf.Module

python
1import tensorflow as tf
2
3class MyModel(tf.Module):
4    def __init__(self, input_size, output_size, name=None):
5        super().__init__(name=name)
6        self.W = tf.Variable(
7            tf.random.normal([input_size, output_size]),
8            name='weights'
9        )
10        self.b = tf.Variable(
11            tf.zeros([output_size]),
12            name='bias'
13        )
14
15    @tf.function
16    def __call__(self, x):
17        return tf.nn.softmax(tf.matmul(x, self.W) + self.b)
18
19# Usage — no sessions needed
20model = MyModel(10, 3)
21data = tf.random.normal([5, 10])
22result = model(data)
23print(result.shape)  # (5, 3)

Using Keras (Most Common)

python
1import tensorflow as tf
2
3class MyModel(tf.keras.Model):
4    def __init__(self, output_size):
5        super().__init__()
6        self.dense1 = tf.keras.layers.Dense(64, activation='relu')
7        self.dense2 = tf.keras.layers.Dense(32, activation='relu')
8        self.output_layer = tf.keras.layers.Dense(output_size, activation='softmax')
9
10    def call(self, inputs, training=False):
11        x = self.dense1(inputs)
12        x = self.dense2(x)
13        return self.output_layer(x)
14
15# Create and use the model
16model = MyModel(output_size=3)
17
18# Forward pass
19data = tf.random.normal([5, 10])
20predictions = model(data)
21print(predictions.shape)  # (5, 3)
22
23# Training
24model.compile(
25    optimizer='adam',
26    loss='sparse_categorical_crossentropy',
27    metrics=['accuracy']
28)
29model.fit(train_data, train_labels, epochs=10)

Using @tf.function for Graph Optimization

@tf.function converts a Python function to a TensorFlow graph for performance:

python
1class Predictor(tf.Module):
2    def __init__(self):
3        super().__init__()
4        self.model = tf.keras.Sequential([
5            tf.keras.layers.Dense(128, activation='relu'),
6            tf.keras.layers.Dense(64, activation='relu'),
7            tf.keras.layers.Dense(10, activation='softmax')
8        ])
9
10    @tf.function(input_signature=[tf.TensorSpec(shape=[None, 784], dtype=tf.float32)])
11    def predict(self, x):
12        return self.model(x, training=False)
13
14    @tf.function
15    def train_step(self, x, y, optimizer, loss_fn):
16        with tf.GradientTape() as tape:
17            predictions = self.model(x, training=True)
18            loss = loss_fn(y, predictions)
19        gradients = tape.gradient(loss, self.model.trainable_variables)
20        optimizer.apply_gradients(zip(gradients, self.model.trainable_variables))
21        return loss
22
23# Usage
24predictor = Predictor()
25result = predictor.predict(tf.random.normal([32, 784]))

Saving and Loading Models

TF2 SavedModel

python
1# Save
2model = MyModel(output_size=3)
3model(tf.random.normal([1, 10]))  # Build the model first
4tf.saved_model.save(model, 'saved_model/')
5
6# Load and use outside the original class
7loaded = tf.saved_model.load('saved_model/')
8result = loaded(tf.random.normal([5, 10]))

Keras Save/Load

python
1# Save weights only
2model.save_weights('model_weights/')
3
4# Save entire model
5model.save('full_model/')
6
7# Load
8loaded_model = tf.keras.models.load_model('full_model/')
9predictions = loaded_model.predict(test_data)

Migrating from TF1 to TF2

TF1 ConceptTF2 Equivalent
tf.Graph()Not needed (eager by default)
tf.Session()Not needed
tf.placeholder()Function arguments
tf.Variable()tf.Variable() (same)
sess.run(op, feed_dict)Direct function call
tf.global_variables_initializer()Variables init on creation
Graph inside a classtf.Module or tf.keras.Model

Common Pitfalls

  • TF1 graph leaks: In TF1, if you create ops without with self.graph.as_default():, they go to the default global graph, causing cross-contamination between model instances. Always scope your ops.
  • Session lifecycle: In TF1 with persistent sessions, forgetting to close the session leaks GPU memory. Use context managers (with tf.Session() as sess:) or call .close() explicitly.
  • @tf.function retracing: In TF2, @tf.function traces a new graph for each unique input signature. Passing Python values (not tensors) causes retracing. Use input_signature to fix the trace.
  • Variable creation in @tf.function: Creating tf.Variable inside @tf.function on the first call works but raises errors on subsequent calls. Create variables in __init__, not in the decorated function.
  • Training vs inference: In Keras, pass training=True during training (enables dropout, batch norm) and training=False during inference. Forgetting this affects model behavior.

Summary

  • TF1 requires explicit graph and session management — create graphs in __init__, run in methods with sessions
  • TF2 uses eager execution — define models as tf.Module or tf.keras.Model subclasses with no session boilerplate
  • Use @tf.function to convert methods to optimized graph operations for performance
  • Prefer tf.keras.Model for full training/saving/loading support
  • When migrating from TF1, replace tf.placeholder with function arguments and remove all tf.Session usage
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
Codemia
© 2026 Codemia
Resources
Legal
Social
All Rights Reserved.