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:
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:
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()
TensorFlow 2.x Approach (Recommended)
TF2 uses eager execution by default. No graphs or sessions to manage:
Using tf.Module
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)
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:
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
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
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 Concept | TF2 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 class | tf.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
