Introduction to TensorFlow Architecture
TensorFlow is one of the most widely used open-source machine learning frameworks in the world. While many developers start by writing simple scripts in Jupyter notebooks, production-grade machine learning systems require a well-thought-out architecture. Understanding TensorFlow's internal design and applying proven design patterns to your project structure is essential for building maintainable, scalable, and reproducible ML systems.
In this tutorial, we will explore the foundational architecture of TensorFlow, examine key design patterns that work well with it, and walk through a recommended project structure that you can adopt for your own ML projects.
Understanding TensorFlow's Core Architecture
At its highest level, TensorFlow is designed around a dataflow programming model. Computation is represented as a directed graph where nodes represent mathematical operations and edges represent the multidimensional data arrays (tensors) flowing between them. This design allows TensorFlow to optimize, parallelize, and distribute computations efficiently across CPUs, GPUs, and TPUs.
Key Architectural Layers
TensorFlow's architecture can be broken down into several distinct layers, each serving a specific purpose:
- Frontend APIs: Python, C++, and other language bindings that developers interact with directly. Keras is now the official high-level API.
- Execution Layer: Handles eager execution (default in TensorFlow 2.x) and graph-based execution via
tf.function. - Graph and Runtime: Converts operations into optimized computation graphs and schedules them across available devices.
- Kernels: Hardware-specific implementations of operations (CPU, GPU, TPU kernels).
- Distributed Runtime: Manages distributed training across multiple workers and parameter servers using
tf.distribute. - Storage and I/O:
tf.datapipelines, TFRecord format, and integration with cloud storage systems.
Eager Execution vs. Graph Execution
TensorFlow 2.x defaults to eager execution, which evaluates operations immediately, making debugging and development much easier. However, for production performance, TensorFlow can convert Python functions into static graphs using the @tf.function decorator.
import tensorflow as tf
# Eager execution (default)
x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
y = tf.constant([[5.0, 6.0], [7.0, 8.0]])
result = tf.matmul(x, y)
print(result.numpy())
# Graph execution for performance
@tf.function
def compute_matrix_product(a, b):
return tf.matmul(a, b)
graph_result = compute_matrix_product(x, y)
print(graph_result.numpy())
Understanding when to use eager execution versus graph execution is a fundamental architectural decision. Use eager execution during development and debugging, and leverage tf.function for performance-critical training loops and inference pipelines.
Why Architecture Matters in TensorFlow Projects
Machine learning projects differ from traditional software projects in several ways. They involve experimentation, data pipelines, model training, evaluation, deployment, and monitoring. Without a clear architecture, projects quickly become tangled messes of notebooks and scripts that are impossible to reproduce or maintain.
Key Benefits of a Well-Structured TensorFlow Project
- Reproducibility: Clear separation of configuration, data, and code ensures experiments can be reproduced.
- Scalability: Modular components can be scaled independently, from data ingestion to distributed training.
- Collaboration: A standardized structure helps teams work together without stepping on each other's toes.
- Deployment Readiness: Clean separation between training and serving code simplifies model deployment.
- Testing: Modular code is testable code. Unit tests for data processing and model logic become straightforward.
Essential Design Patterns for TensorFlow
Design patterns provide reusable solutions to common problems. In TensorFlow projects, several patterns have emerged as best practices for organizing code and managing complexity.
1. The Model-Builder Pattern
Instead of defining models inline, encapsulate model construction in builder functions or classes. This makes it easy to swap architectures, tune hyperparameters, and test different configurations.
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
class ModelBuilder:
"""Builder class for constructing neural network models."""
@staticmethod
def build_cnn(num_classes: int, input_shape: tuple = (224, 224, 3)) -> keras.Model:
inputs = keras.Input(shape=input_shape)
x = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(inputs)
x = layers.MaxPooling2D((2, 2))(x)
x = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(x)
x = layers.MaxPooling2D((2, 2))(x)
x = layers.Conv2D(128, (3, 3), activation='relu', padding='same')(x)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dense(256, activation='relu')(x)
x = layers.Dropout(0.5)(x)
outputs = layers.Dense(num_classes, activation='softmax')(x)
model = keras.Model(inputs=inputs, outputs=outputs, name='cnn_classifier')
return model
@staticmethod
def build_mlp(num_classes: int, input_dim: int = 128) -> keras.Model:
inputs = keras.Input(shape=(input_dim,))
x = layers.Dense(256, activation='relu')(inputs)
x = layers.Dropout(0.3)(x)
x = layers.Dense(128, activation='relu')(x)
x = layers.Dropout(0.3)(x)
outputs = layers.Dense(num_classes, activation='softmax')(x)
model = keras.Model(inputs=inputs, outputs=outputs, name='mlp_classifier')
return model
# Usage
model = ModelBuilder.build_cnn(num_classes=10)
model.summary()
2. The Configuration Pattern
Separate configuration from code using dataclasses or configuration files (YAML, JSON). This allows you to run experiments with different hyperparameters without modifying source code.
from dataclasses import dataclass, field
from typing import List, Optional
import yaml
@dataclass
class DataConfig:
data_dir: str = "data/raw"
batch_size: int = 32
image_size: tuple = (224, 224)
num_classes: int = 10
validation_split: float = 0.2
augment: bool = True
@dataclass
class ModelConfig:
architecture: str = "cnn"
learning_rate: float = 0.001
optimizer: str = "adam"
dropout_rate: float = 0.5
fine_tune_layers: int = 0
@dataclass
class TrainConfig:
epochs: int = 50
early_stopping_patience: int = 5
checkpoint_dir: str = "checkpoints"
log_dir: str = "logs"
mixed_precision: bool = True
@dataclass
class ExperimentConfig:
experiment_name: str = "baseline"
data: DataConfig = field(default_factory=DataConfig)
model: ModelConfig = field(default_factory=ModelConfig)
train: TrainConfig = field(default_factory=TrainConfig)
@classmethod
def from_yaml(cls, path: str) -> "ExperimentConfig":
with open(path, 'r') as f:
config_dict = yaml.safe_load(f)
return cls(
experiment_name=config_dict.get('experiment_name', 'baseline'),
data=DataConfig(**config_dict.get('data', {})),
model=ModelConfig(**config_dict.get('model', {})),
train=TrainConfig(**config_dict.get('train', {})),
)
# Usage
config = ExperimentConfig.from_yaml("configs/experiment_01.yaml")
print(f"Running experiment: {config.experiment_name}")
print(f"Batch size: {config.data.batch_size}")
print(f"Learning rate: {config.model.learning_rate}")
3. The Data Pipeline Pattern
Use tf.data to build efficient, reusable data pipelines. Encapsulate data loading, preprocessing, and augmentation in dedicated classes or functions. This pattern ensures consistent data handling between training and evaluation.
import tensorflow as tf
class DataPipeline:
"""Encapsulates data loading, preprocessing, and augmentation."""
def __init__(self, config):
self.config = config
self.autotune = tf.data.AUTOTUNE
def load_dataset(self, data_dir: str) -> tf.data.Dataset:
"""Load images from directory structure."""
dataset = tf.keras.utils.image_dataset_from_directory(
data_dir,
image_size=self.config.image_size,
batch_size=self.config.batch_size,
shuffle=True,
seed=42,
validation_split=self.config.validation_split,
subset='training',
)
return dataset
def augment_data(self, image, label):
"""Apply data augmentation to training samples."""
image = tf.image.random_flip_left_right(image)
image = tf.image.random_brightness(image, max_delta=0.2)
image = tf.image.random_contrast(image, lower=0.8, upper=1.2)
image = tf.image.random_saturation(image, lower=0.8, upper=1.2)
image = tf.clip_by_value(image, 0.0, 255.0)
return image, label
def normalize(self, image, label):
"""Normalize pixel values to [0, 1]."""
image = tf.cast(image, tf.float32) / 255.0
return image, label
def build_train_pipeline(self, data_dir: str) -> tf.data.Dataset:
dataset = self.load_dataset(data_dir)
if self.config.augment:
dataset = dataset.map(self.augment_data, num_parallel_calls=self.autotune)
dataset = dataset.map(self.normalize, num_parallel_calls=self.autotune)
dataset = dataset.prefetch(self.autotune)
return dataset
def build_eval_pipeline(self, data_dir: str) -> tf.data.Dataset:
dataset = tf.keras.utils.image_dataset_from_directory(
data_dir,
image_size=self.config.image_size,
batch_size=self.config.batch_size,
shuffle=False,
seed=42,
validation_split=self.config.validation_split,
subset='validation',
)
dataset = dataset.map(self.normalize, num_parallel_calls=self.autotune)
dataset = dataset.prefetch(self.autotune)
return dataset
# Usage
pipeline = DataPipeline(config.data)
train_ds = pipeline.build_train_pipeline("data/raw")
val_ds = pipeline.build_eval_pipeline("data/raw")
4. The Trainer Pattern
Encapsulate the training logic in a dedicated trainer class. This separates the training loop from model definition and data preparation, making the codebase cleaner and more maintainable.
import tensorflow as tf
import os
from datetime import datetime
class Trainer:
"""Handles model training, checkpointing, and logging."""
def __init__(self, model, config, train_ds, val_ds):
self.model = model
self.config = config
self.train_ds = train_ds
self.val_ds = val_ds
# Configure mixed precision if enabled
if config.mixed_precision:
tf.keras.mixed_precision.set_global_policy('mixed_float16')
self._setup_callbacks()
self._compile_model()
def _compile_model(self):
optimizer_map = {
'adam': tf.keras.optimizers.Adam(learning_rate=self.config.learning_rate),
'sgd': tf.keras.optimizers.SGD(learning_rate=self.config.learning_rate, momentum=0.9),
'adamw': tf.keras.optimizers.AdamW(learning_rate=self.config.learning_rate),
}
optimizer = optimizer_map.get(self.config.optimizer,
tf.keras.optimizers.Adam(self.config.learning_rate))
self.model.compile(
optimizer=optimizer,
loss='sparse_categorical_crossentropy',
metrics=['accuracy'],
)
def _setup_callbacks(self):
timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
self.callbacks = [
tf.keras.callbacks.ModelCheckpoint(
filepath=os.path.join(self.config.checkpoint_dir, 'best_model.keras'),
monitor='val_accuracy',
save_best_only=True,
mode='max',
verbose=1,
),
tf.keras.callbacks.EarlyStopping(
monitor='val_loss',
patience=self.config.early_stopping_patience,
restore_best_weights=True,
verbose=1,
),
tf.keras.callbacks.TensorBoard(
log_dir=os.path.join(self.config.log_dir, timestamp),
histogram_freq=1,
),
tf.keras.callbacks.ReduceLROnPlateau(
monitor='val_loss',
factor=0.5,
patience=3,
min_lr=1e-7,
verbose=1,
),
]
def train(self):
"""Execute the training loop."""
print("Starting training...")
history = self.model.fit(
self.train_ds,
validation_data=self.val_ds,
epochs=self.config.epochs,
callbacks=self.callbacks,
)
print("Training complete.")
return history
def evaluate(self, test_ds):
"""Evaluate the model on test data."""
results = self.model.evaluate(test_ds, return_dicts=True)
print(f"Test results: {results}")
return results
# Usage
trainer = Trainer(model, config.train, train_ds, val_ds)
history = trainer.train()
5. The Strategy Pattern for Distributed Training
TensorFlow provides tf.distribute.Strategy for distributing training across multiple GPUs or workers. Wrap your model creation and training logic inside a strategy scope to make it distribution-ready.
import tensorflow as tf
class DistributedTrainer:
"""Trainer that supports distributed training across multiple devices."""
def __init__(self, config):
self.config = config
# Automatically detect available GPUs
num_gpus = len(tf.config.list_physical_devices('GPU'))
if num_gpus > 1:
self.strategy = tf.distribute.MirroredStrategy()
print(f"Using MirroredStrategy with {num_gpus} GPUs")
elif num_gpus == 1:
self.strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
print("Using single GPU")
else:
self.strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0")
print("Using CPU")
def build_and_train(self, model_fn, train_ds, val_ds):
with self.strategy.scope():
model = model_fn()
model.compile(
optimizer=tf.keras.optimizers.Adam(self.config.learning_rate),
loss='sparse_categorical_crossentropy',
metrics=['accuracy'],
)
# Dataset is automatically distributed
history = model.fit(
train_ds,
validation_data=val_ds,
epochs=self.config.epochs,
)
return model, history
# Usage
dist_trainer = DistributedTrainer(config.train)
model, history = dist_trainer.build_and_train(
model_fn=lambda: ModelBuilder.build_cnn(config.data.num_classes),
train_ds=train_ds,
val_ds=val_ds,
)
6. The Registry Pattern for Custom Components
When working with many model variants, loss functions, or metrics, a registry pattern keeps things organized. It allows you to register components by name and retrieve them dynamically, which is especially useful for configuration-driven experiments.
class Registry:
"""Generic registry for mapping names to classes/functions."""
def __init__(self, name: str):
self.name = name
self._registry = {}
def register(self, name: str):
def decorator(cls):
self._registry[name] = cls
return cls
return decorator
def get(self, name: str):
if name not in self._registry:
raise KeyError(f"'{name}' not found in {self.name} registry. "
f"Available: {list(self._registry.keys())}")
return self._registry[name]
def list_available(self):
return list(self._registry.keys())
# Create registries
MODEL_REGISTRY = Registry("models")
LOSS_REGISTRY = Registry("losses")
# Register models
@MODEL_REGISTRY.register("simple_cnn")
class SimpleCNN(tf.keras.Model):
def __init__(self, num_classes=10):
super().__init__()
self.conv1 = tf.keras.layers.Conv2D(32, 3, activation='relu')
self.pool = tf.keras.layers.GlobalAveragePooling2D()
self.fc = tf.keras.layers.Dense(num_classes, activation='softmax')
def call(self, inputs):
x = self.conv1(inputs)
x = self.pool(x)
return self.fc(x)
@MODEL_REGISTRY.register("resnet_transfer")
class ResNetTransfer(tf.keras.Model):
def __init__(self, num_classes=10):
super().__init__()
self.base = tf.keras.applications.ResNet50(
weights='imagenet', include_top=False, input_shape=(224, 224, 3)
)
self.base.trainable = False
self.gap = tf.keras.layers.GlobalAveragePooling2D()
self.classifier = tf.keras.layers.Dense(num_classes, activation='softmax')
def call(self, inputs):
x = self.base(inputs)
x = self.gap(x)
return self.classifier(x)
# Retrieve model by name from config
model_class = MODEL_REGISTRY.get(config.model.architecture)
model = model_class(num_classes=config.data.num_classes)
print(f"Available models: {MODEL_REGISTRY.list_available()}")
Recommended Project Structure
Now that we have covered the key design patterns, let us put them together into a cohesive project structure. The following layout is inspired by best practices from the ML engineering community and works well for small to medium-sized TensorFlow projects.
my_tensorflow_project/
├── configs/
│ ├── experiment_01.yaml
│ ├── experiment_02.yaml
│ └── production.yaml
├── data/
│ ├── raw/
│ ├── processed/
│ ├── external/
│ └── interim/
├── models/
│ ├── checkpoints/
│ ├── saved_models/
│ └── exports/
├── notebooks/
│ ├── exploration.ipynb
│ └── visualization.ipynb
├── src/
│ ├── __init__.py
│ ├── data/
│ │ ├── __init__.py
│ │ ├── dataset.py
│ │ ├── preprocessing.py
│ │ └── augmentation.py
│ ├── models/
│ │ ├── __init__.py
│ │ ├── base.py
│ │ ├── cnn.py
│ │ ├── resnet.py
│ │ └── registry.py
│ ├── training/
│ │ ├── __init__.py
│ │ ├── trainer.py
│ │ ├── callbacks.py
│ │ └── distributed.py
│ ├── evaluation/
│ │ ├── __init__.py
│ │ ├── metrics.py
│ │ └── evaluator.py
│ ├── serving/
│ │ ├── __init__.py
│ │ ├── inference.py
│ │ └── serving_client.py
│ └── utils/
│ ├── __init__.py
│ ├── config.py
│ ├── logger.py
│ └── seed.py
├── tests/
│ ├── __init__.py
│ ├── test_data.py
│ ├── test_models.py
│ └── test_training.py
├── scripts/
│ ├── train.py
│ ├── evaluate.py
│ ├── export_model.py
│ └── download_data.py
├── .gitignore
├── requirements.txt
├── setup.py
├── Makefile
└── README.md
Explanation of Each Directory
- configs/: YAML configuration files for different experiments and environments. Keeping configs separate from code enables easy experimentation.
- data/: Follows the Cookiecutter Data Science convention.
raw/contains immutable original data,processed/holds cleaned and transformed data ready for training. - models/: Stores model artifacts including checkpoints, saved models, and exported versions for serving.
- notebooks/: Jupyter notebooks for exploration and visualization only. Production code should never live here.
- src/: The main source code, organized by concern: data, models, training, evaluation, serving, and utilities.
- tests/: Unit and integration tests for all components.
- scripts/: Entry-point scripts that tie together the modules in
src/for specific tasks like training or evaluation.
Example Training Script
Here is how the scripts/train.py entry point ties everything together using the patterns we discussed:
"""Training entry point script."""
import argparse
import sys
import os
# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))
from utils.config import ExperimentConfig
from data.dataset import DataPipeline
from models.registry import MODEL_REGISTRY
from training.trainer import Trainer
from utils.logger import setup_logger
from utils.seed import set_seed
def main():
parser = argparse.ArgumentParser(description="Train a TensorFlow model")
parser.add_argument('--config', type=str, required=True,
help='Path to YAML config file')
parser.add_argument('--data-dir', type=str, default=None,
help='Override data directory')
args = parser.parse_args()
# Load configuration
config = ExperimentConfig.from_yaml(args.config)
if args.data_dir:
config.data.data_dir = args.data_dir
# Setup logging and seeding
logger = setup_logger(config.experiment_name)
set_seed(42)
logger.info(f"Starting experiment: {config.experiment_name}")
# Build data pipeline
pipeline = DataPipeline(config.data)
train_ds = pipeline.build_train_pipeline(config.data.data_dir)
val_ds = pipeline.build_eval_pipeline(config.data.data_dir)
logger.info("Data pipelines built successfully.")
# Build model from registry
model_class = MODEL_REGISTRY.get(config.model.architecture)
model = model_class(num_classes=config.data.num_classes)
logger.info(f"Model '{config.model.architecture}' built.")
# Train
trainer = Trainer(model, config.train, train_ds, val_ds)
history = trainer.train()
# Save final model
save_path = os.path.join('models', 'saved_models',
f"{config.experiment_name}.keras")
model.save(save_path)
logger.info(f"Model saved to {save_path}")
print("Training pipeline completed successfully!")
if __name__ == '__main__':
main()
Example Configuration File
Here is a sample YAML configuration file that the training script consumes:
# configs/experiment_01.yaml
experiment_name: "cnn_baseline_v1"
data:
data_dir: "data/raw"
batch_size: 64
image_size: [224, 224]
num_classes: 10
validation_split: 0.2
augment: true
model:
architecture: "simple_cnn"
learning_rate: 0.001
optimizer: "adam"
dropout_rate: 0.5
fine_tune_layers: 0
train:
epochs: 100
early_stopping_patience: 10
checkpoint_dir: "models/checkpoints"
log_dir: "logs"
mixed_precision: true
Best Practices for TensorFlow Architecture
1. Embrace Modularity
Every component should have a single responsibility. Data loading, preprocessing, model definition, training, and evaluation should each live in their own modules. This makes it easy to test, replace, and extend individual pieces without affecting the rest of the system.
2. Use Configuration-Driven Development
Never hardcode hyperparameters in your source code. Use YAML or JSON configuration files and dataclasses to manage experiment settings. This enables reproducibility and makes it trivial to run multiple experiments with different parameters.
3. Optimize Data Pipelines Early
Data loading is often the bottleneck in TensorFlow training. Always use tf.data with prefetch, cache, and parallel map operations. Here is a quick optimization checklist:
# Optimized tf.data pipeline
dataset = (
tf.data.Dataset.from_tensor_slices((images, labels))
.shuffle(buffer_size=10000)
.map(preprocess_fn, num_parallel_calls=tf.data.AUTOTUNE)
.batch(batch_size)
.cache() # Cache after preprocessing to avoid recomputation
.prefetch(tf.data.AUTOTUNE) # Overlap data loading with training
)
4. Version Your Models and Data
Use clear naming conventions for saved models that include the experiment name, date, and key metrics. Consider using TensorFlow's SavedModel format for deployment and track model versions using MLflow, Weights & Biases, or a similar tool.
import os
from datetime import datetime
def save_model_with_metadata(model, config, metrics):
"""Save model with a descriptive name and metadata."""
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
val_acc = metrics.get('val_accuracy', 0.0)
model_name = f"{config.experiment_name}_valacc{val_acc:.4f}_{timestamp}"
save_path = os.path.join('models', 'saved_models', model_name)
model.save(save_path)
# Save metadata alongside
metadata = {
'experiment_name': config.experiment_name,
'architecture': config.model.architecture,
'val_accuracy': val_acc,
'timestamp': timestamp,
'config': vars(config),
}
import json
with open(os.path.join(save_path, 'metadata.json'), 'w') as f:
json.dump(metadata, f, indent=2)
print(f"Model saved to: {save_path}")
return save_path
5. Write Tests for Critical Components
Test your data pipelines, model output shapes, and custom metrics. TensorFlow models can be tested by checking output shapes and verifying that gradients flow correctly.
import tensorflow as tf
import pytest
class TestModelBuilder:
def test_cnn_output_shape(self):
model = ModelBuilder.build_cnn(num_classes=10, input_shape=(64, 64, 3))
dummy_input = tf.random.normal((1, 64, 64, 3))
output = model(dummy_input)
assert output.shape == (1, 10), f"Expected (1, 10), got {output.shape}"
def test_cnn_output_sum_to_one(self):
model = ModelBuilder.build_cnn(num_classes=5, input_shape=(64, 64, 3))
dummy_input = tf.random.normal((1, 64, 64, 3))
output = model(dummy_input)
assert tf.reduce_allclose(tf.reduce_sum(output, axis=-1), 1.0, atol=1e-5)
def test_model_is_trainable(self):
model = ModelBuilder.build_cnn(num_classes=3, input_shape=(32, 32, 3))
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
assert len(model.trainable_variables) > 0
class TestDataPipeline:
def test_normalization_range(self):
config = DataConfig(image_size=(32, 32), batch_size=1)
pipeline = DataPipeline(config)
dummy_image = tf.random.uniform((32, 32, 3), minval=0, maxval=255)
normalized, _ = pipeline.normalize(dummy_image, label=0)
assert tf.reduce_max(normalized) <= 1.0
assert tf.reduce_min(normalized) >= 0.0
6. Leverage Mixed Precision Training
On modern GPUs (Volta and later), mixed precision training can significantly speed up training and reduce memory usage with minimal impact on accuracy.
# Enable mixed precision globally
tf.keras.mixed_precision.set_global_policy('mixed_float16')
# Build and train your model as usual
model = ModelBuilder.build_cnn(num_classes=10)
# Ensure the output layer uses float32 for numerical stability
# This is handled automatically by Keras when using mixed_float16 policy
model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
loss='sparse_categorical_crossentropy',
metrics=['accuracy'],
)
7. Plan for Serving from the Start
Design your project so that the transition from training to serving is seamless. Export models in the SavedModel format and write a simple inference module that can be used by a REST API or gRPC service.
import tensorflow as tf
import numpy as np
class ModelServer:
"""Simple inference wrapper for serving a trained model."""
def __init__(self, model_path: str):
self.model = tf.keras.models.load_model(model_path)
self.input_shape = self.model.input_shape[1:]
def preprocess(self, image: np.ndarray) -> np.ndarray:
"""Preprocess raw image for inference."""
image = tf.image.resize(image, self.input_shape[:2])
image = tf.cast(image, tf.float32) / 255.0
image = tf.expand_dims(image, axis=0)
return image.numpy()
def predict(self, image: np.ndarray) -> dict:
"""Run inference and return predictions."""
processed = self.preprocess(image)
predictions = self.model.predict(processed, verbose=0)
predicted_class = int(np.argmax(predictions[0]))
confidence = float(predictions[0][predicted_class])
return {
'predicted_class': predicted_class,
'confidence': confidence,
'all_probabilities': predictions[0].tolist(),
}
# Usage
server = ModelServer("models/saved_models/cnn_baseline_v1.keras")
dummy_image = np.random.randint(0, 255, (300, 300, 3), dtype=np.uint8)
result = server.predict(dummy_image)
print(f"Prediction: {result}")
8. Use Logging Instead of Print Statements
In production code, use Python's logging module instead of print(). This allows you to control log levels, format output, and direct logs to files or monitoring systems.
import logging
import os
from datetime import datetime
def setup_logger(experiment_name: str, log_dir: str = "logs") -> logging.Logger:
"""Configure and return a logger for the experiment."""
os.makedirs(log_dir, exist_ok=True)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
log_file = os.path.join(log_dir, f"{experiment_name}_{timestamp}.log")
logger = logging.getLogger(experiment_name)
logger.setLevel(logging.DEBUG)
# File handler
file_handler = logging.FileHandler(log_file)
file_handler.setLevel(logging.DEBUG)
file_handler.setFormatter(
logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
)
# Console handler
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
console_handler.setFormatter(
logging.Formatter('%(levelname)s - %(message)s')
)
logger.addHandler(file_handler)
logger.addHandler(console_handler)
return logger
Putting It All Together
Let us look at a complete, minimal example that demonstrates how all these patterns work together. This example shows the full flow from configuration to training to saving a model.
"""Complete example: putting all patterns together."""
import tensorflow as tf
from dataclasses import dataclass, field
# --- Configuration ---
@dataclass
class Config:
data_dir: str = "data/raw"
batch_size: int = 32
image_size: tuple = (128, 128)
num_classes: int = 5
learning_rate: float = 0.001
epochs: int = 20
# --- Data Pipeline ---
def build_data_pipeline(config: Config) -> tuple:
train_ds = tf.keras.utils.image_dataset_from_directory(
config.data_dir,
image_size=config.image_size,
batch_size=config.batch_size,
validation_split=0.2,
subset='training',
seed=42,
)
val_ds = tf.keras.utils.image_dataset_from_directory(
config.data_dir,
image_size=config.image_size,
batch_size=config.batch_size,
validation_split=0.2,
subset='validation',
seed=42,
)
normalization = tf.keras.layers.Rescaling(1.0 / 255)
train_ds = train_ds.map(lambda x, y: (normalization(x), y)).prefetch(tf.data.AUTOTUNE)
val_ds = val_ds.map(lambda x, y: (normalization(x), y)).prefetch(tf.data.AUTOTUNE)
return train_ds, val_ds
# --- Model Builder ---
def build_model(config: Config) -> tf.keras.Model:
model = tf.keras.Sequential([
tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(*config.image_size, 3)),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Conv2D(64, 3, activation='relu'),
tf.keras.layers.MaxPooling2D(),
tf.keras.layers.Conv2D(128, 3, activation='relu'),
tf.keras.layers.GlobalAveragePooling2D(),
tf.keras.layers.Dense(128, activation='relu'),
tf.keras.layers.Dropout(0.5),
tf.keras.layers.Dense(config.num_classes, activation='softmax'),
])
model.compile(
optimizer=tf.keras.optimizers.Adam(config.learning_rate),
loss='sparse_categorical_crossentropy',
metrics=['accuracy'],
)
return model
# --- Training ---
def train(config: Config):
print(f"=== Starting Training ===")
print(f"Data dir: {config.data_dir}")
print(f"Batch size: {config.batch_size}")
print(f"Epochs: {config.epochs}")
train_ds, val_ds = build_data_pipeline(config)
model = build_model(config)
callbacks = [
tf.keras.callbacks.EarlyStopping(patience=5, restore_best_weights=True),
tf.keras.callbacks.ModelCheckpoint('models/best_model.keras', save_best_only=True),
]
history = model.fit(
train_ds,
validation_data=val_ds,
epochs=config.epochs,
callbacks=callbacks,
)
# Save final model
model.save('models/final_model.keras')
print("=== Training Complete ===")
print(f"Best val accuracy: {max(history.history['val_accuracy']):.4f}")
return model, history
# --- Run ---
if __name__ == '__main__':
config = Config(
data_dir="data/raw",
batch_size=32,
image_size=(128, 128),
num_classes=5,
learning_rate=0.001,
epochs=20,
)
model, history = train(config)
Conclusion
Building production-grade TensorFlow applications requires more than just knowing the API. By understanding TensorFlow's underlying architecture and applying proven design patterns such as the model-builder, configuration, data pipeline, trainer, strategy, and registry patterns, you can create ML systems that are modular, maintainable, and scalable. A well-organized project structure with clear separation