Introduction

TensorFlow code often starts as a short script and then grows into a training pipeline, evaluation workflow, and model-serving path. Using object-oriented structure helps keep those concerns separate so the code remains testable and predictable when the project stops being a notebook experiment.

Model the Responsibilities First

The main OOP decision is not whether to create classes, but what each class should own. A clean TensorFlow design usually separates:

• model architecture
• training behavior
• evaluation and metrics
• persistence or checkpointing
• configuration

If one class loads data, builds the network, trains, evaluates, logs, and saves artifacts, the class is already doing too much. Composition is usually a better fit than inheritance for ML code because responsibilities change independently.

Wrap the Keras Model in a Domain Class

A common and practical approach is to make one class responsible for creating and compiling the network.

1import tensorflow as tf
2
3
4class ClassifierModel:
5    def __init__(self, input_dim: int, hidden_units: int, num_classes: int):
6        self.model = tf.keras.Sequential(
7            [
8                tf.keras.layers.Input(shape=(input_dim,)),
9                tf.keras.layers.Dense(hidden_units, activation="relu"),
10                tf.keras.layers.Dense(num_classes, activation="softmax"),
11            ]
12        )
13
14    def compile(self, learning_rate: float = 1e-3) -> None:
15        self.model.compile(
16            optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
17            loss="sparse_categorical_crossentropy",
18            metrics=["accuracy"],
19        )

This class owns architecture and compile choices. It does not own training loops or storage.

Move Training Logic into a Trainer Class

Training policy changes often. You may add early stopping, checkpoints, class weights, mixed precision, or custom callbacks. Those changes should not force architecture edits.

1class Trainer:
2    def __init__(self, classifier: ClassifierModel):
3        self.classifier = classifier
4
5    def fit(self, x_train, y_train, x_val, y_val, epochs: int = 10):
6        callbacks = [
7            tf.keras.callbacks.EarlyStopping(
8                monitor="val_loss",
9                patience=2,
10                restore_best_weights=True,
11            )
12        ]
13        return self.classifier.model.fit(
14            x_train,
15            y_train,
16            validation_data=(x_val, y_val),
17            epochs=epochs,
18            callbacks=callbacks,
19            verbose=2,
20        )

That separation makes it easier to reuse the same model class with different training policies.

Add Evaluation and Persistence Boundaries

Evaluation and persistence should be explicit too. That keeps file-system behavior and metric reporting out of the training flow.

1class Evaluator:
2    def __init__(self, classifier: ClassifierModel):
3        self.classifier = classifier
4
5    def evaluate(self, x_test, y_test):
6        return self.classifier.model.evaluate(x_test, y_test, verbose=0)
7
8
9class ModelStore:
10    @staticmethod
11    def save(classifier: ClassifierModel, path: str) -> None:
12        classifier.model.save(path)
13
14    @staticmethod
15    def load(path: str):
16        return tf.keras.models.load_model(path)

This is especially useful once model saving rules differ across environments such as local training, CI, and deployment packaging.

Full Runnable Example

Putting the classes together gives a small but maintainable training flow.

1import numpy as np
2
3x = np.random.rand(600, 8).astype("float32")
4y = np.random.randint(0, 3, size=(600,))
5
6x_train, x_val = x[:500], x[500:]
7y_train, y_val = y[:500], y[500:]
8
9classifier = ClassifierModel(input_dim=8, hidden_units=32, num_classes=3)
10classifier.compile()
11
12trainer = Trainer(classifier)
13trainer.fit(x_train, y_train, x_val, y_val, epochs=3)
14
15evaluator = Evaluator(classifier)
16loss, acc = evaluator.evaluate(x_val, y_val)
17print("val_loss", loss)
18print("val_acc", acc)

This pattern is still simple enough for small projects, but it also scales better than one procedural script.

Know When to Subclass tf.keras.Model

Not every OOP TensorFlow design needs a custom subclass of tf.keras.Model. Use subclassing when you need custom forward logic, multiple inputs with special handling, or a custom train_step. For ordinary tabular or feed-forward models, wrapping a Sequential or functional Keras model is often cleaner and easier for teammates to read.

Use the least complex abstraction that supports the current requirement.

Common Pitfalls

• Building one giant class that owns data loading, training, evaluation, and persistence together.
• Using inheritance where simple composition would keep responsibilities clearer.
• Hiding hyperparameters in module-level globals instead of passing configuration explicitly.
• Writing custom model subclasses before the project actually needs custom forward or training logic.
• Treating “OOP style” as a goal by itself instead of using it to improve boundaries and testability.

Summary

• OOP TensorFlow code should separate architecture, training, evaluation, and persistence.
• Composition is usually a better default than deep inheritance.
• A wrapper around a Keras model is often enough for clean design.
• Trainer and evaluator classes keep policy changes out of the model definition.
• Use custom tf.keras.Model subclassing only when the model behavior truly demands it.