Part 9: Deployment, Performance & Best Practices

SavedModel Format

The SavedModel is TensorFlow's universal serialization format for trained models. Unlike simple weight checkpoints, a SavedModel captures the complete computation graph, variables, assets (vocabulary files, etc.), and serving signatures — everything needed to run inference without the original training code. It's the standard format accepted by TensorFlow Serving, TFLite, TensorFlow.js, and TensorFlow Hub.

A SavedModel directory contains:

• saved_model.pb — serialized computation graph (protocol buffer)
• variables/ — trained weights (variable values)
• assets/ — external files referenced by the graph (vocabularies, lookup tables)
• fingerprint.pb — model fingerprint for integrity checks

import tensorflow as tf
import numpy as np

# Build a simple model to demonstrate SavedModel export
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# Train on synthetic data
X_train = np.random.randn(1000, 10).astype(np.float32)
y_train = (X_train[:, 0] > 0).astype(np.float32)
model.fit(X_train, y_train, epochs=5, batch_size=32, verbose=0)

# Save as SavedModel (directory format)
export_path = '/tmp/my_model/1'  # Version subdirectory convention
model.save(export_path)
print(f"Model saved to: {export_path}")

# Inspect the saved directory structure
import os
for root, dirs, files in os.walk(export_path):
    level = root.replace(export_path, '').count(os.sep)
    indent = ' ' * 2 * level
    print(f'{indent}{os.path.basename(root)}/')
    subindent = ' ' * 2 * (level + 1)
    for file in files:
        filepath = os.path.join(root, file)
        size_kb = os.path.getsize(filepath) / 1024
        print(f'{subindent}{file} ({size_kb:.1f} KB)')

# Load the model back (no training code needed)
loaded_model = tf.keras.models.load_model(export_path)
test_input = np.random.randn(5, 10).astype(np.float32)
predictions = loaded_model.predict(test_input, verbose=0)
print(f"\nPredictions shape: {predictions.shape}")
print(f"Sample predictions: {predictions.flatten()[:3]}")

Signatures & Serving Functions

For serving, you often want explicit signatures — named input/output mappings that define the model's API contract. You can define custom serving functions using @tf.function with input_signature for strict type and shape enforcement:

import tensorflow as tf
import numpy as np

# Build and train a model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(3, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
X = np.random.randn(500, 10).astype(np.float32)
y = np.random.randint(0, 3, 500)
model.fit(X, y, epochs=3, verbose=0)

# Define custom serving function with explicit signature
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 10], dtype=tf.float32)])
def serve(input_tensor):
    """Custom serving function with preprocessing."""
    # Normalize input (example preprocessing)
    normalized = (input_tensor - tf.reduce_mean(input_tensor, axis=1, keepdims=True))
    predictions = model(normalized, training=False)
    # Return class probabilities and predicted class
    return {
        'probabilities': predictions,
        'predicted_class': tf.argmax(predictions, axis=1)
    }

# Save with custom signatures
export_path = '/tmp/custom_serving_model/1'
tf.saved_model.save(
    model,
    export_path,
    signatures={'serving_default': serve}
)

# Load and inspect signatures
loaded = tf.saved_model.load(export_path)
print("Available signatures:", list(loaded.signatures.keys()))

# Use the serving function
infer = loaded.signatures['serving_default']
test = tf.constant(np.random.randn(3, 10).astype(np.float32))
result = infer(test)
print(f"Output keys: {list(result.keys())}")
print(f"Probabilities: {result['probabilities'].numpy()}")
print(f"Predicted classes: {result['predicted_class'].numpy()}")

TensorFlow Serving

TensorFlow Serving is a high-performance serving system designed for production environments. It provides model versioning (hot-swapping new model versions without downtime), batching (combining multiple requests for GPU efficiency), and both REST and gRPC APIs. The standard deployment uses Docker containers for isolation and reproducibility.

flowchart LR
    A[Train Model] --> B[Export SavedModel]
    B --> C[Version Directory
/models/v1, /models/v2]
    C --> D[TF Serving Container]
    D --> E[REST API :8501]
    D --> F[gRPC API :8500]
    E --> G[Client Applications]
    F --> G
    C --> H[Model Config
Versioning Policy]
    H --> D
    I[Monitoring] --> D
    D --> J[Metrics & Logs]
            

# Pull TensorFlow Serving Docker image
docker pull tensorflow/serving:latest

# Run TF Serving with a model directory
# Model must be in /models/model_name/version_number/ format
docker run -d --name tf_serving \
  -p 8501:8501 \
  -p 8500:8500 \
  -v "/path/to/models:/models" \
  -e MODEL_NAME=my_model \
  tensorflow/serving:latest

# With batching enabled (batching_parameters.txt)
cat > /tmp/batching_config.txt << 'EOF'
max_batch_size { value: 32 }
batch_timeout_micros { value: 5000 }
max_enqueued_batches { value: 100 }
num_batch_threads { value: 4 }
EOF

docker run -d --name tf_serving_batched \
  -p 8501:8501 \
  -v "/path/to/models:/models" \
  -e MODEL_NAME=my_model \
  tensorflow/serving:latest \
  --enable_batching=true \
  --batching_parameters_file=/models/batching_config.txt

# Health check
curl http://localhost:8501/v1/models/my_model

# Check model status and metadata
curl http://localhost:8501/v1/models/my_model/metadata

REST & gRPC APIs

TF Serving exposes two APIs: REST (HTTP/JSON on port 8501) for simplicity and broad compatibility, and gRPC (port 8500) for maximum throughput with binary serialization. For latency-sensitive production workloads, gRPC is typically 2-5× faster due to Protocol Buffer encoding and HTTP/2 streaming.

import numpy as np
import json
import requests

# REST API inference request
# Assumes TF Serving running on localhost:8501 with model "my_model"
test_data = np.random.randn(3, 10).tolist()

# REST prediction request
payload = json.dumps({
    "signature_name": "serving_default",
    "instances": test_data
})

headers = {"Content-Type": "application/json"}
response = requests.post(
    "http://localhost:8501/v1/models/my_model:predict",
    data=payload,
    headers=headers
)

print(f"Status: {response.status_code}")
predictions = response.json()
print(f"Predictions: {predictions}")

# gRPC API (faster for production)
# pip install tensorflow-serving-api grpcio
import grpc
import tensorflow as tf
from tensorflow_serving.apis import predict_pb2, prediction_service_pb2_grpc

# Create gRPC channel
channel = grpc.insecure_channel('localhost:8500')
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)

# Build request
request = predict_pb2.PredictRequest()
request.model_spec.name = 'my_model'
request.model_spec.signature_name = 'serving_default'

# Convert numpy array to TensorProto
input_data = np.random.randn(3, 10).astype(np.float32)
request.inputs['input_1'].CopyFrom(
    tf.make_tensor_proto(input_data, shape=input_data.shape)
)

# Send request
result = stub.Predict(request, timeout=10.0)
output = tf.make_ndarray(result.outputs['output_0'])
print(f"gRPC predictions shape: {output.shape}")
print(f"gRPC predictions: {output}")

TensorFlow Lite

TensorFlow Lite (TFLite) optimizes models for mobile phones, embedded devices, and edge hardware. It converts TensorFlow models into a compact FlatBuffer format (.tflite) with an interpreter optimized for ARM CPUs, GPUs, and specialized accelerators (Edge TPU, Hexagon DSP). The key tool is quantization — reducing weight precision from float32 to int8, shrinking model size by 4× and often improving inference speed by 2-3×.

The quantization formula converts floating-point values to integers:

$$q = \text{round}\left(\frac{x}{\text{scale}}\right) + \text{zero\_point}$$

where $\text{scale}$ and $\text{zero\_point}$ are calibrated from the data distribution to minimize quantization error.

import tensorflow as tf
import numpy as np

# Build and train a model for edge deployment
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(16, 3, activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Conv2D(32, 3, activation='relu'),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Train on synthetic MNIST-like data
X_train = np.random.randn(1000, 28, 28, 1).astype(np.float32)
y_train = np.random.randint(0, 10, 1000)
model.fit(X_train, y_train, epochs=3, verbose=0)

# --- Dynamic Range Quantization (post-training, easiest) ---
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_dynamic = converter.convert()
print(f"Dynamic quantized model size: {len(tflite_dynamic) / 1024:.1f} KB")

# --- Full Integer Quantization (requires representative dataset) ---
def representative_dataset():
    """Provide sample inputs for calibration."""
    for _ in range(100):
        yield [np.random.randn(1, 28, 28, 1).astype(np.float32)]

converter_int = tf.lite.TFLiteConverter.from_keras_model(model)
converter_int.optimizations = [tf.lite.Optimize.DEFAULT]
converter_int.representative_dataset = representative_dataset
converter_int.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter_int.inference_input_type = tf.int8
converter_int.inference_output_type = tf.int8
tflite_int8 = converter_int.convert()
print(f"Full int8 model size: {len(tflite_int8) / 1024:.1f} KB")

# --- Float16 Quantization (good GPU acceleration) ---
converter_fp16 = tf.lite.TFLiteConverter.from_keras_model(model)
converter_fp16.optimizations = [tf.lite.Optimize.DEFAULT]
converter_fp16.target_spec.supported_types = [tf.float16]
tflite_fp16 = converter_fp16.convert()
print(f"Float16 model size: {len(tflite_fp16) / 1024:.1f} KB")

# Compare with original
original_size = sum(np.prod(w.shape) * 4 for w in model.get_weights()) / 1024
print(f"\nOriginal model weights: {original_size:.1f} KB")
print(f"Compression ratios:")
print(f"  Dynamic: {original_size / (len(tflite_dynamic)/1024):.1f}x")
print(f"  Int8:    {original_size / (len(tflite_int8)/1024):.1f}x")
print(f"  FP16:    {original_size / (len(tflite_fp16)/1024):.1f}x")

Running TFLite Inference

Here is the implementation for Running TFLite Inference. Each code example below is self-contained and can be run independently:

import tensorflow as tf
import numpy as np

# Load a TFLite model and run inference
# (Using in-memory model from conversion; in practice, load from .tflite file)
model = tf.keras.Sequential([
    tf.keras.layers.Dense(32, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(3, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
X = np.random.randn(100, 10).astype(np.float32)
y = np.random.randint(0, 3, 100)
model.fit(X, y, epochs=2, verbose=0)

# Convert to TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

# Save to file
with open('/tmp/model.tflite', 'wb') as f:
    f.write(tflite_model)

# Load and run with TFLite interpreter
interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

# Get input/output details
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(f"Input shape: {input_details[0]['shape']}")
print(f"Input dtype: {input_details[0]['dtype']}")
print(f"Output shape: {output_details[0]['shape']}")

# Run inference on a single sample
test_input = np.random.randn(1, 10).astype(np.float32)
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
output = interpreter.get_tensor(output_details[0]['index'])
print(f"\nPrediction: {output}")
print(f"Predicted class: {np.argmax(output)}")

# Benchmark inference time
import time
num_runs = 1000
start = time.time()
for _ in range(num_runs):
    interpreter.set_tensor(input_details[0]['index'], test_input)
    interpreter.invoke()
elapsed = time.time() - start
print(f"\nAverage inference time: {elapsed/num_runs*1000:.2f} ms")

TensorFlow.js

TensorFlow.js brings ML to the browser and Node.js. Models run entirely client-side — no server round-trips, instant inference, and complete data privacy. The WebGL backend leverages GPU acceleration for near-native performance on most operations. You can convert existing Python models or train directly in JavaScript.

# Convert a SavedModel to TensorFlow.js format
# pip install tensorflowjs

# From SavedModel directory
tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_format=tfjs_graph_model \
    --signature_name=serving_default \
    /tmp/my_model/1 \
    /tmp/tfjs_model/

# From Keras H5 file
tensorflowjs_converter \
    --input_format=keras \
    /tmp/model.h5 \
    /tmp/tfjs_model/

# With quantization for smaller download
tensorflowjs_converter \
    --input_format=tf_saved_model \
    --output_format=tfjs_graph_model \
    --quantize_uint8 \
    /tmp/my_model/1 \
    /tmp/tfjs_model_quantized/

Browser Inference & Transfer Learning

Here is the implementation for Browser Inference & Transfer Learning. Each code example below is self-contained and can be run independently:

// TensorFlow.js browser inference
// Include: <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>

async function loadAndPredict() {
    // Load converted model
    const model = await tf.loadGraphModel('/models/tfjs_model/model.json');

    // Create input tensor (e.g., preprocessed image)
    const inputTensor = tf.randomNormal([1, 224, 224, 3]);

    // Run inference
    const predictions = model.predict(inputTensor);
    const results = await predictions.data();
    console.log('Predictions:', results);

    // Clean up tensors to prevent memory leaks
    inputTensor.dispose();
    predictions.dispose();
}

// Transfer learning in the browser
async function transferLearning() {
    // Load pretrained MobileNet (feature extractor)
    const mobilenet = await tf.loadLayersModel(
        'https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json'
    );

    // Freeze base model layers
    for (const layer of mobilenet.layers) {
        layer.trainable = false;
    }

    // Get intermediate layer output as features
    const featureLayer = mobilenet.getLayer('conv_pw_13_relu');
    const truncatedModel = tf.model({
        inputs: mobilenet.inputs,
        outputs: featureLayer.output
    });

    // Add custom classification head
    const model = tf.sequential();
    model.add(tf.layers.inputLayer({ inputShape: featureLayer.outputShape.slice(1) }));
    model.add(tf.layers.globalAveragePooling2d({}));
    model.add(tf.layers.dense({ units: 64, activation: 'relu' }));
    model.add(tf.layers.dropout({ rate: 0.3 }));
    model.add(tf.layers.dense({ units: 5, activation: 'softmax' }));

    model.compile({
        optimizer: tf.train.adam(0.001),
        loss: 'categoricalCrossentropy',
        metrics: ['accuracy']
    });

    console.log('Transfer learning model ready for fine-tuning');
    model.summary();
}

loadAndPredict();
transferLearning();

Multi-GPU Training

Training on multiple GPUs can reduce training time linearly (ideally). TensorFlow's Distribution Strategies API provides a clean abstraction — wrap your model creation and training in a strategy scope, and TensorFlow handles gradient synchronization, data sharding, and variable mirroring automatically.

flowchart TD
    A[Distribution Strategies] --> B[MirroredStrategy]
    A --> C[MultiWorkerMirroredStrategy]
    A --> D[TPUStrategy]
    A --> E[ParameterServerStrategy]

    B --> B1[Single machine
multiple GPUs]
    B --> B2[Synchronous
all-reduce]

    C --> C1[Multiple machines
multiple GPUs]
    C --> C2[Synchronous
ring all-reduce]

    D --> D1[Google Cloud TPU
pods]
    D --> D2[Replicated across
TPU cores]

    E --> E1[Large-scale
async training]
    E --> E2[Separate parameter
server processes]
            

import tensorflow as tf
import numpy as np

# --- MirroredStrategy: single machine, multiple GPUs ---
# Replicates model on each GPU, synchronizes gradients via all-reduce
strategy = tf.distribute.MirroredStrategy()
print(f"Number of devices: {strategy.num_replicas_in_sync}")

# Everything inside the scope is distributed
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(256, activation='relu', input_shape=(100,)),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dropout(0.3),
        tf.keras.layers.Dense(128, activation='relu'),
        tf.keras.layers.BatchNormalization(),
        tf.keras.layers.Dense(10, activation='softmax')
    ])

    # Optimizer and compilation inside strategy scope
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.001),
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

# Scale batch size with number of GPUs
# Each GPU processes (global_batch_size / num_gpus) samples
GLOBAL_BATCH_SIZE = 256 * strategy.num_replicas_in_sync
print(f"Global batch size: {GLOBAL_BATCH_SIZE}")
print(f"Per-replica batch size: {GLOBAL_BATCH_SIZE // strategy.num_replicas_in_sync}")

# Create dataset (automatically sharded across replicas)
X_train = np.random.randn(10000, 100).astype(np.float32)
y_train = np.random.randint(0, 10, 10000)

dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train))
dataset = dataset.shuffle(10000).batch(GLOBAL_BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

# Training is automatically distributed
history = model.fit(dataset, epochs=5, verbose=1)
print(f"Final accuracy: {history.history['accuracy'][-1]:.4f}")

Multi-Worker & TPU Strategy

Here is the implementation for Multi-Worker & TPU Strategy. Each code example below is self-contained and can be run independently:

import tensorflow as tf
import numpy as np
import json
import os

# --- MultiWorkerMirroredStrategy ---
# For training across multiple machines
# Each worker needs TF_CONFIG environment variable set

# Example TF_CONFIG for worker 0 of 2
tf_config = {
    "cluster": {
        "worker": ["worker0:12345", "worker1:12345"]
    },
    "task": {"type": "worker", "index": 0}
}
os.environ['TF_CONFIG'] = json.dumps(tf_config)

# Strategy automatically discovers cluster from TF_CONFIG
multi_worker_strategy = tf.distribute.MultiWorkerMirroredStrategy()

with multi_worker_strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(128, activation='relu', input_shape=(50,)),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# --- TPUStrategy ---
# For Google Cloud TPU training
try:
    resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
    tf.config.experimental_connect_to_cluster(resolver)
    tf.tpu.experimental.initialize_tpu_system(resolver)
    tpu_strategy = tf.distribute.TPUStrategy(resolver)
    print(f"TPU cores: {tpu_strategy.num_replicas_in_sync}")
except ValueError:
    print("No TPU detected — falling back to default strategy")
    tpu_strategy = tf.distribute.get_strategy()

# Batch size scaling rule of thumb:
# Linear scaling: lr_new = lr_base * num_gpus
# Warmup: gradually increase lr over first few epochs
base_lr = 0.001
num_replicas = 4  # Example: 4 GPUs
scaled_lr = base_lr * num_replicas
warmup_epochs = 5

print(f"\nBatch size scaling:")
print(f"  Base LR: {base_lr}")
print(f"  Scaled LR ({num_replicas} GPUs): {scaled_lr}")
print(f"  Warmup epochs: {warmup_epochs}")

XLA Compilation

XLA (Accelerated Linear Algebra) is a domain-specific compiler that optimizes TensorFlow computations. It fuses multiple operations into single optimized kernels, eliminates intermediate memory allocations, and applies algebraic simplifications. The result: fewer kernel launches, better memory locality, and often 20-50% speedup for compute-bound models.

XLA excels at:

• Kernel fusion — combining element-wise operations (add, multiply, activation) into single GPU kernels
• Memory optimization — eliminating temporary buffers for intermediate results
• Layout optimization — choosing optimal tensor memory layouts for the target hardware
• Constant folding — pre-computing values known at compile time

import tensorflow as tf
import numpy as np
import time

# --- Enable XLA with jit_compile=True ---
# Method 1: On individual tf.functions
@tf.function(jit_compile=True)
def xla_matmul(a, b):
    """XLA-compiled matrix multiplication with activation."""
    result = tf.matmul(a, b)
    result = tf.nn.relu(result)
    result = result + tf.ones_like(result)  # Fused with relu by XLA
    return tf.nn.softmax(result, axis=-1)

# Test XLA function
a = tf.random.normal([512, 256])
b = tf.random.normal([256, 128])

# Warmup (first call triggers compilation)
_ = xla_matmul(a, b)

# Benchmark
start = time.time()
for _ in range(1000):
    _ = xla_matmul(a, b)
tf.debugging.assert_shapes([])  # Sync GPU
xla_time = time.time() - start
print(f"XLA compiled: {xla_time:.3f}s for 1000 iterations")

# Compare without XLA
@tf.function(jit_compile=False)
def no_xla_matmul(a, b):
    result = tf.matmul(a, b)
    result = tf.nn.relu(result)
    result = result + tf.ones_like(result)
    return tf.nn.softmax(result, axis=-1)

_ = no_xla_matmul(a, b)  # Warmup
start = time.time()
for _ in range(1000):
    _ = no_xla_matmul(a, b)
no_xla_time = time.time() - start
print(f"Without XLA:  {no_xla_time:.3f}s for 1000 iterations")
print(f"Speedup: {no_xla_time/xla_time:.2f}x")

# Method 2: Compile entire model training step
model = tf.keras.Sequential([
    tf.keras.layers.Dense(256, activation='relu', input_shape=(100,)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# jit_compile on the model itself
model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy'],
    jit_compile=True  # Compile training step with XLA
)

X = np.random.randn(5000, 100).astype(np.float32)
y = np.random.randint(0, 10, 5000)
model.fit(X, y, epochs=3, batch_size=64, verbose=1)
print("\nXLA-compiled training complete!")

When XLA Helps vs. Hurts

Mixed Precision

Mixed precision training uses float16 for forward/backward computation while maintaining float32 master copies of weights. This halves memory usage for activations and enables Tensor Core acceleration on NVIDIA GPUs (Volta+), delivering 2-3× speedup with negligible accuracy loss. The key insight is that most neural network operations are robust to reduced precision, but weight updates require full precision to accumulate small gradients.

import tensorflow as tf
import numpy as np

# Enable mixed precision globally
tf.keras.mixed_precision.set_global_policy('mixed_float16')
policy = tf.keras.mixed_precision.global_policy()
print(f"Policy: {policy.name}")
print(f"Compute dtype: {policy.compute_dtype}")  # float16
print(f"Variable dtype: {policy.variable_dtype}")  # float32

# Build model — layers automatically use float16 compute, float32 variables
model = tf.keras.Sequential([
    tf.keras.layers.Dense(512, activation='relu', input_shape=(100,)),
    tf.keras.layers.Dense(256, activation='relu'),
    tf.keras.layers.Dense(128, activation='relu'),
    # CRITICAL: Final layer must output float32 for numerical stability
    tf.keras.layers.Dense(10, dtype='float32')  # Explicit float32 output
])

# Use loss scaling to prevent float16 underflow in gradients
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

model.compile(
    optimizer=optimizer,
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    metrics=['accuracy']
)

# Verify dtypes
for layer in model.layers:
    print(f"{layer.name}: compute={layer.dtype_policy.compute_dtype}, "
          f"variable={layer.dtype_policy.variable_dtype}")

# Train normally — mixed precision is transparent
X_train = np.random.randn(5000, 100).astype(np.float32)
y_train = np.random.randint(0, 10, 5000)

history = model.fit(X_train, y_train, epochs=5, batch_size=128, verbose=1)
print(f"\nFinal accuracy: {history.history['accuracy'][-1]:.4f}")

# Memory savings: activations stored in float16 (half the memory)
# Allows ~2x larger batch sizes on same GPU
print("\nMixed precision benefits:")
print("  - 2x smaller activation memory")
print("  - 2-3x faster on Tensor Cores (V100, A100, RTX 3000+)")
print("  - Negligible accuracy impact for most models")

# Reset to default after demonstration
tf.keras.mixed_precision.set_global_policy('float32')

Loss Scaling

When gradients are very small (common in deep networks), float16 can underflow to zero. Loss scaling multiplies the loss by a large factor before backpropagation (keeping gradients in float16 representable range), then divides gradients by the same factor before the weight update. TensorFlow's mixed_float16 policy handles this automatically with dynamic loss scaling.

Profiling & Optimization

Knowing where your training time goes is essential for optimization. TensorFlow's profiler identifies whether you're bottlenecked by input pipeline (CPU-bound data loading), compute (GPU operations), or host-device communication (data transfers). The TensorBoard profiler provides visual timelines, op-level breakdowns, and actionable recommendations.

import tensorflow as tf
import numpy as np

# --- TensorBoard Profiler Setup ---
# Create a log directory for profiling
log_dir = '/tmp/tf_profile_logs'

# Method 1: Profile with TensorBoard callback
model = tf.keras.Sequential([
    tf.keras.layers.Dense(256, activation='relu', input_shape=(100,)),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# Profile steps 10-20 (skip warmup steps)
tensorboard_callback = tf.keras.callbacks.TensorBoard(
    log_dir=log_dir,
    profile_batch='10,20'  # Profile batch range
)

X = np.random.randn(5000, 100).astype(np.float32)
y = np.random.randint(0, 10, 5000)
model.fit(X, y, epochs=2, batch_size=64, callbacks=[tensorboard_callback], verbose=0)
print(f"Profile saved to: {log_dir}")
print("View with: tensorboard --logdir /tmp/tf_profile_logs")

# Method 2: Programmatic profiling with tf.profiler
tf.profiler.experimental.start(log_dir + '/manual_profile')

# Run the operations you want to profile
for _ in range(50):
    batch_x = tf.random.normal([64, 100])
    with tf.GradientTape() as tape:
        predictions = model(batch_x, training=True)
        loss = tf.reduce_mean(predictions)
    gradients = tape.gradient(loss, model.trainable_variables)

tf.profiler.experimental.stop()
print("Manual profiling complete")

# Method 3: Trace individual operations
@tf.function
def train_step(x, y):
    with tf.GradientTape() as tape:
        predictions = model(x, training=True)
        loss = tf.keras.losses.sparse_categorical_crossentropy(y, predictions)
        loss = tf.reduce_mean(loss)
    gradients = tape.gradient(loss, model.trainable_variables)
    model.optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

# Use tf.summary.trace for operation-level detail
tf.summary.trace_on(graph=True, profiler=True)
sample_x = tf.random.normal([64, 100])
sample_y = tf.constant(np.random.randint(0, 10, 64))
train_step(sample_x, sample_y)

with tf.summary.create_file_writer(log_dir + '/trace').as_default():
    tf.summary.trace_export(name="train_step", step=0)
print("Trace exported — view Graph tab in TensorBoard")

Identifying Bottlenecks

Model Interpretability

Understanding why a model makes specific predictions is critical for debugging, building trust, and meeting regulatory requirements. TensorFlow supports several interpretability techniques: Grad-CAM (visual explanations for CNNs), Integrated Gradients (attribution for any differentiable model), and SHAP (game-theoretic feature importance). These help answer "which input features drove this prediction?"

import tensorflow as tf
import numpy as np

# --- Grad-CAM: Visual Explanations for CNNs ---
# Highlights which regions of an image the model "looks at"

# Build a simple CNN
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(32, 32, 3)),
    tf.keras.layers.Conv2D(64, 3, activation='relu', name='last_conv'),
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(10, activation='softmax')
])
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# Train briefly
X = np.random.randn(500, 32, 32, 3).astype(np.float32)
y = np.random.randint(0, 10, 500)
model.fit(X, y, epochs=2, verbose=0)

def grad_cam(model, image, last_conv_layer_name, pred_index=None):
    """Generate Grad-CAM heatmap for a given image."""
    # Create a model that outputs both convolution output and predictions
    grad_model = tf.keras.Model(
        inputs=model.input,
        outputs=[
            model.get_layer(last_conv_layer_name).output,
            model.output
        ]
    )

    # Compute gradients of predicted class w.r.t. conv output
    with tf.GradientTape() as tape:
        conv_output, predictions = grad_model(image[np.newaxis])
        if pred_index is None:
            pred_index = tf.argmax(predictions[0])
        class_output = predictions[:, pred_index]

    # Gradient of the predicted class w.r.t. feature map
    grads = tape.gradient(class_output, conv_output)

    # Global average pooling of gradients → channel importance weights
    weights = tf.reduce_mean(grads, axis=(1, 2))  # (1, num_channels)

    # Weighted combination of feature maps
    cam = tf.reduce_sum(conv_output[0] * weights[0], axis=-1)

    # ReLU and normalize
    cam = tf.maximum(cam, 0)
    cam = cam / (tf.reduce_max(cam) + 1e-8)
    return cam.numpy()

# Generate Grad-CAM for a sample image
sample_image = np.random.randn(32, 32, 3).astype(np.float32)
heatmap = grad_cam(model, sample_image, 'last_conv')
print(f"Grad-CAM heatmap shape: {heatmap.shape}")
print(f"Heatmap range: [{heatmap.min():.3f}, {heatmap.max():.3f}]")
print("Overlay this on the original image to see which regions drive the prediction")

Integrated Gradients

Here is the implementation for Integrated Gradients. Each code example below is self-contained and can be run independently:

import tensorflow as tf
import numpy as np

# --- Integrated Gradients ---
# Model-agnostic attribution that satisfies key axioms
# (sensitivity, implementation invariance)

# Simple model for demonstration
model = tf.keras.Sequential([
    tf.keras.layers.Dense(64, activation='relu', input_shape=(20,)),
    tf.keras.layers.Dense(32, activation='relu'),
    tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy')
X = np.random.randn(500, 20).astype(np.float32)
y = (X[:, 0] + X[:, 1] > 0).astype(np.float32)  # Features 0,1 are important
model.fit(X, y, epochs=10, verbose=0)

def integrated_gradients(model, input_sample, baseline=None, steps=50):
    """
    Compute Integrated Gradients attribution.

    Accumulates gradients along the path from baseline to input,
    attributing the prediction difference to each input feature.
    """
    if baseline is None:
        baseline = tf.zeros_like(input_sample)

    # Generate interpolated inputs along straight-line path
    alphas = tf.linspace(0.0, 1.0, steps + 1)
    interpolated = tf.stack([
        baseline + alpha * (input_sample - baseline)
        for alpha in alphas
    ])

    # Compute gradients at each interpolation point
    with tf.GradientTape() as tape:
        tape.watch(interpolated)
        predictions = model(interpolated)

    gradients = tape.gradient(predictions, interpolated)

    # Approximate integral using trapezoidal rule
    avg_gradients = tf.reduce_mean(gradients, axis=0)

    # Scale by (input - baseline)
    integrated_grads = avg_gradients * (input_sample - baseline)
    return integrated_grads.numpy()

# Compute attributions
sample = X[0:1].astype(np.float32)
sample_tensor = tf.constant(sample)
attributions = integrated_gradients(model, sample_tensor[0])

# Display feature importance
print("Integrated Gradients attributions (top 5 features):")
importance = np.abs(attributions[0])
top_features = np.argsort(importance)[::-1][:5]
for idx in top_features:
    print(f"  Feature {idx}: {attributions[0][idx]:+.4f} (|attr|={importance[idx]:.4f})")

# Verify: features 0 and 1 should have highest attributions
print(f"\nFeature 0 rank: {list(top_features).index(0) + 1 if 0 in top_features else '>5'}")
print(f"Feature 1 rank: {list(top_features).index(1) + 1 if 1 in top_features else '>5'}")
print("(Features 0 and 1 are the true signal — others are noise)")

Production Best Practices

Moving from experimentation to production requires systematic processes around model versioning, monitoring, testing, and continuous retraining. Google's internal experience (codified in the "ML Test Score" paper) shows that ML code is typically <5% of a production ML system — the rest is data pipelines, monitoring, infrastructure, and configuration management.

import tensorflow as tf
import numpy as np
import json
import time
from datetime import datetime

# --- Model Versioning & Registry Pattern ---
class ModelRegistry:
    """Simple model registry for versioning and metadata tracking."""

    def __init__(self, base_path='/tmp/model_registry'):
        self.base_path = base_path
        self.metadata_path = f"{base_path}/registry.json"

    def register_model(self, model, model_name, metrics, description=""):
        """Save model with version number and metadata."""
        import os
        os.makedirs(self.base_path, exist_ok=True)

        # Load existing registry
        try:
            with open(self.metadata_path, 'r') as f:
                registry = json.load(f)
        except (FileNotFoundError, json.JSONDecodeError):
            registry = {"models": {}}

        # Determine version
        if model_name not in registry["models"]:
            registry["models"][model_name] = []
        version = len(registry["models"][model_name]) + 1

        # Save model
        export_path = f"{self.base_path}/{model_name}/v{version}"
        model.save(export_path)

        # Record metadata
        metadata = {
            "version": version,
            "path": export_path,
            "timestamp": datetime.now().isoformat(),
            "metrics": metrics,
            "description": description,
            "status": "staged"  # staged → canary → production → archived
        }
        registry["models"][model_name].append(metadata)

        with open(self.metadata_path, 'w') as f:
            json.dump(registry, f, indent=2)

        print(f"Registered {model_name} v{version}")
        print(f"  Metrics: {metrics}")
        print(f"  Status: staged")
        return version

    def promote(self, model_name, version, new_status):
        """Promote model to new status (canary/production/archived)."""
        with open(self.metadata_path, 'r') as f:
            registry = json.load(f)

        models = registry["models"][model_name]
        models[version - 1]["status"] = new_status
        print(f"Promoted {model_name} v{version} → {new_status}")

        with open(self.metadata_path, 'w') as f:
            json.dump(registry, f, indent=2)

# Demo: register a model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(32, activation='relu', input_shape=(10,)),
    tf.keras.layers.Dense(1, activation='sigmoid')
])
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
X = np.random.randn(200, 10).astype(np.float32)
y = (X[:, 0] > 0).astype(np.float32)
model.fit(X, y, epochs=5, verbose=0)

registry = ModelRegistry()
registry.register_model(
    model,
    model_name="fraud_detector",
    metrics={"accuracy": 0.92, "auc": 0.96, "f1": 0.89},
    description="Baseline logistic model with dense layers"
)

TFX Pipeline Overview

TFX (TensorFlow Extended) is Google's end-to-end ML platform for production pipelines. It provides standardized components for each stage of the ML lifecycle, all orchestrated by Apache Beam, Airflow, or Kubeflow Pipelines:

import tensorflow as tf
import numpy as np

# --- TFX Pipeline Components Overview ---
# TFX provides a production ML pipeline framework

# Component 1: ExampleGen — Ingests data into the pipeline
# Supports CSV, TFRecord, BigQuery, custom formats
print("TFX Pipeline Components:")
print("=" * 50)

components = {
    "ExampleGen": "Ingest & split data (train/eval/test)",
    "StatisticsGen": "Compute dataset statistics (TFDV)",
    "SchemaGen": "Infer data schema from statistics",
    "ExampleValidator": "Detect anomalies, drift, skew",
    "Transform": "Feature engineering (tf.Transform)",
    "Trainer": "Train model (Keras/Estimator)",
    "Tuner": "Hyperparameter search (KerasTuner)",
    "Evaluator": "Validate model quality (TFMA)",
    "InfraValidator": "Test model serves correctly",
    "Pusher": "Deploy to TF Serving/TFLite/TFJS"
}

for i, (component, description) in enumerate(components.items(), 1):
    print(f"  {i:2d}. {component:20s} → {description}")

# --- Monitoring for Data/Concept Drift ---
# Simple drift detection using statistical tests
def detect_drift(reference_data, production_data, threshold=0.1):
    """
    Detect distribution drift using Population Stability Index (PSI).
    PSI > 0.1 suggests moderate drift; PSI > 0.25 indicates significant drift.
    """
    def compute_psi(expected, actual, bins=10):
        # Bin the data
        breakpoints = np.percentile(expected, np.linspace(0, 100, bins + 1))
        breakpoints[0] = -np.inf
        breakpoints[-1] = np.inf

        expected_counts = np.histogram(expected, bins=breakpoints)[0] / len(expected)
        actual_counts = np.histogram(actual, bins=breakpoints)[0] / len(actual)

        # Avoid division by zero
        expected_counts = np.maximum(expected_counts, 1e-4)
        actual_counts = np.maximum(actual_counts, 1e-4)

        # PSI formula
        psi = np.sum((actual_counts - expected_counts) *
                     np.log(actual_counts / expected_counts))
        return psi

    # Check each feature
    drift_scores = []
    for col in range(reference_data.shape[1]):
        psi = compute_psi(reference_data[:, col], production_data[:, col])
        drift_scores.append(psi)

    avg_psi = np.mean(drift_scores)
    max_psi = np.max(drift_scores)

    alert = "🚨 SIGNIFICANT" if max_psi > 0.25 else "⚠️ MODERATE" if max_psi > threshold else "✅ STABLE"
    return {"avg_psi": avg_psi, "max_psi": max_psi, "alert": alert, "per_feature": drift_scores}

# Simulate drift detection
np.random.seed(42)
reference = np.random.randn(10000, 5)  # Training data distribution
production_stable = np.random.randn(1000, 5)  # No drift
production_drifted = np.random.randn(1000, 5) + 0.5  # Shifted distribution

print("\n\nDrift Detection Results:")
print("-" * 40)
result_stable = detect_drift(reference, production_stable)
print(f"Stable data:  {result_stable['alert']} (max PSI={result_stable['max_psi']:.4f})")

result_drifted = detect_drift(reference, production_drifted)
print(f"Drifted data: {result_drifted['alert']} (max PSI={result_drifted['max_psi']:.4f})")