AI & Machine Learning
Model Training with TensorFlow: Data Prep, Training & Evaluation
TensorFlow and Keras are industry-standard tools for developing, training, and deploying deep learning models. Managing the entire training process — from data preprocessing to hyperparameter optimization — is critical for a successful AI project. In this article, we will examine every stage of model training with detailed code examples.
Building the Data Pipeline
Efficient Data Loading with tf.data
When working with large datasets, efficiently feeding data to the GPU is critical for performance. The tf.data API optimizes this process:
import tensorflow as tf
import numpy as np
# Data loading function
def load_and_preprocess(image_path, label):
"""Load image from disk and preprocess."""
image = tf.io.read_file(image_path)
image = tf.image.decode_jpeg(image, channels=3)
image = tf.image.resize(image, [224, 224])
image = tf.cast(image, tf.float32) / 255.0
# Normalize with ImageNet mean
mean = tf.constant([0.485, 0.456, 0.406])
std = tf.constant([0.229, 0.224, 0.225])
image = (image - mean) / std
return image, label
# Data augmentation function
def augment(image, label):
"""Apply random transformations for training data."""
image = tf.image.random_flip_left_right(image)
image = tf.image.random_brightness(image, 0.2)
image = tf.image.random_contrast(image, 0.8, 1.2)
image = tf.image.random_saturation(image, 0.8, 1.2)
return image, label
# High-performance data pipeline
def build_pipeline(file_paths, labels, batch_size=32, training=True):
"""Create an optimized tf.data pipeline."""
ds = tf.data.Dataset.from_tensor_slices((file_paths, labels))
if training:
ds = ds.shuffle(buffer_size=10000, seed=42)
ds = ds.map(load_and_preprocess, num_parallel_calls=tf.data.AUTOTUNE)
if training:
ds = ds.map(augment, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.batch(batch_size)
ds = ds.prefetch(tf.data.AUTOTUNE)
return dsModel Creation and Compilation
Defining Models with the Keras Functional API
For complex architectures, the Functional API offers a more flexible structure:
from tensorflow.keras import layers, models, optimizers
def build_model(num_classes, learning_rate=0.001):
"""Build a transfer learning classification model."""
# Pre-trained backbone
backbone = tf.keras.applications.EfficientNetV2S(
include_top=False,
weights='imagenet',
input_shape=(224, 224, 3)
)
backbone.trainable = False # Freeze backbone
# Model architecture
inputs = layers.Input(shape=(224, 224, 3))
x = backbone(inputs, training=False)
x = layers.GlobalAveragePooling2D()(x)
x = layers.BatchNormalization()(x)
x = layers.Dense(512, activation='relu')(x)
x = layers.Dropout(0.4)(x)
x = layers.Dense(256, activation='relu')(x)
x = layers.Dropout(0.3)(x)
outputs = layers.Dense(num_classes, activation='softmax')(x)
model = models.Model(inputs=inputs, outputs=outputs)
# Compile model
model.compile(
optimizer=optimizers.Adam(learning_rate=learning_rate),
loss='sparse_categorical_crossentropy',
metrics=[
'accuracy',
tf.keras.metrics.TopKCategoricalAccuracy(k=5, name='top5_accuracy')
]
)
return model
model = build_model(num_classes=100)
model.summary()Training Control with Callbacks
Keras callbacks are powerful tools for automatically managing the training process:
import os
from datetime import datetime
# Log directory
log_dir = os.path.join("logs", datetime.now().strftime("%Y%m%d-%H%M%S"))
# Callback list
callbacks = [
# Early stopping: stop if val_loss doesn't improve for 7 epochs
tf.keras.callbacks.EarlyStopping(
monitor='val_loss',
patience=7,
restore_best_weights=True,
verbose=1
),
# Reduce learning rate: halve if val_loss stalls for 3 epochs
tf.keras.callbacks.ReduceLROnPlateau(
monitor='val_loss',
factor=0.5,
patience=3,
min_lr=1e-7,
verbose=1
),
# Save the best model
tf.keras.callbacks.ModelCheckpoint(
filepath='best_model.keras',
monitor='val_accuracy',
save_best_only=True,
verbose=1
),
# TensorBoard logs
tf.keras.callbacks.TensorBoard(
log_dir=log_dir,
histogram_freq=1,
write_graph=True
),
# CSV logger
tf.keras.callbacks.CSVLogger('training_log.csv')
]Training Loop
Two-Phase Training Strategy
A common strategy in transfer learning projects is to first train only the top layers, then fine-tune:
# Phase 1: Train classification layers only
print("Phase 1: Classification layers training")
history_1 = model.fit(
train_ds,
epochs=20,
validation_data=val_ds,
callbacks=callbacks,
verbose=1
)
# Phase 2: Fine-tuning — unfreeze the last layers of the backbone
print("\nPhase 2: Fine-tuning")
backbone = model.layers[1]
backbone.trainable = True
# Freeze all but the last 30 layers
for layer in backbone.layers[:-30]:
layer.trainable = False
# Recompile with a lower learning rate
model.compile(
optimizer=optimizers.Adam(learning_rate=1e-5),
loss='sparse_categorical_crossentropy',
metrics=['accuracy']
)
history_2 = model.fit(
train_ds,
epochs=30,
validation_data=val_ds,
callbacks=callbacks,
verbose=1
)Custom Training Loop with GradientTape
While Keras's model.fit() method is sufficient for most cases, some scenarios require full control over the training loop. Custom training loops are ideal for non-standard training procedures such as GANs, meta-learning, and research projects. tf.GradientTape provides this control.
GradientTape records the operations performed during the forward pass and then uses these records to automatically compute gradients. This mechanism allows you to differentiate any TensorFlow operation.
# Custom training loop
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy()
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
val_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()
@tf.function # Compile to graph mode for performance
def train_step(images, labels):
"""Single training step."""
with tf.GradientTape() as tape:
predictions = model(images, training=True)
loss = loss_fn(labels, predictions)
# Add L2 regularization loss
total_loss = loss + tf.add_n(model.losses) if model.losses else loss
# Compute and apply gradients
gradients = tape.gradient(total_loss, model.trainable_variables)
# Gradient clipping to prevent instability
gradients = [tf.clip_by_norm(g, 1.0) for g in gradients]
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
train_accuracy.update_state(labels, predictions)
return total_loss
@tf.function
def val_step(images, labels):
"""Single validation step."""
predictions = model(images, training=False)
loss = loss_fn(labels, predictions)
val_accuracy.update_state(labels, predictions)
return loss
# Training loop
NUM_EPOCHS = 30
best_accuracy = 0.0
for epoch in range(NUM_EPOCHS):
train_accuracy.reset_state()
val_accuracy.reset_state()
# Training
for batch_images, batch_labels in train_ds:
loss = train_step(batch_images, batch_labels)
# Validation
for batch_images, batch_labels in val_ds:
val_step(batch_images, batch_labels)
print(f"Epoch {epoch+1}/{NUM_EPOCHS} - "
f"Train Accuracy: {train_accuracy.result():.4f} - "
f"Val Accuracy: {val_accuracy.result():.4f}")
# Save best model
if val_accuracy.result() > best_accuracy:
best_accuracy = val_accuracy.result()
model.save_weights('best_weights.weights.h5')Mixed Precision Training
Mixed precision training significantly speeds up training by performing some computations in FP16 (half precision). Modern GPUs with NVIDIA Tensor Cores execute FP16 operations much faster than FP32. This technique also reduces memory usage, allowing you to work with larger batch sizes.
The key point is that gradients and model weights are still maintained in FP32; only the forward and backward pass computations are done in FP16. TensorFlow manages these transitions automatically.
# Enable mixed precision training
tf.keras.mixed_precision.set_global_policy('mixed_float16')
# When building the model, note: softmax should stay in FP32
inputs = layers.Input(shape=(224, 224, 3))
x = backbone(inputs, training=False)
x = layers.GlobalAveragePooling2D()(x)
x = layers.Dense(512, activation='relu')(x)
x = layers.Dropout(0.4)(x)
# Keep the final layer in FP32 for numerical stability
outputs = layers.Dense(num_classes, activation='softmax', dtype='float32')(x)
model_mp = models.Model(inputs=inputs, outputs=outputs)
# Use optimizer with loss scaling
optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
model_mp.compile(
optimizer=optimizer,
loss='sparse_categorical_crossentropy',
metrics=['accuracy']
)
print(f"Compute policy: {tf.keras.mixed_precision.global_policy()}")Evaluation and Analysis
Visualizing Model Performance
import matplotlib.pyplot as plt
def plot_training(history):
"""Plot training and validation metrics."""
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# Accuracy plot
axes[0].plot(history.history['accuracy'], label='Training')
axes[0].plot(history.history['val_accuracy'], label='Validation')
axes[0].set_title('Model Accuracy')
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Accuracy')
axes[0].legend()
axes[0].grid(True)
# Loss plot
axes[1].plot(history.history['loss'], label='Training')
axes[1].plot(history.history['val_loss'], label='Validation')
axes[1].set_title('Model Loss')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('Loss')
axes[1].legend()
axes[1].grid(True)
plt.tight_layout()
plt.savefig('training_plots.png', dpi=150)
plt.show()Hyperparameter Optimization
Critical hyperparameters for successful model training:
- Learning Rate: The single most important hyperparameter. Too high = divergence, too low = slow convergence
- Batch Size: Larger batch = more stable gradients but more memory. Typical values: 16, 32, 64
- Number of Epochs: Use early stopping for automatic determination
- Optimizer: Adam works well in most cases; SGD + momentum is preferred for fine-tuning
- Weight Decay: Try values between 1e-4 and 1e-2 for L2 regularization
- Dropout Rate: Typically between 0.2 and 0.5
Monitoring with TensorBoard
TensorBoard is a powerful visualization tool for real-time training monitoring:
tensorboard --logdir logs --port 6006What you can monitor with TensorBoard:
- Training and validation metrics (loss, accuracy)
- Layer weight distributions and histograms
- Model graph structure
- Learning rate changes
- Embedding visualizations
Model Quantization for Mobile (TFLite)
Running trained models on mobile devices or embedded systems requires reducing model size and increasing inference speed. TensorFlow Lite (TFLite) provides model quantization for this purpose. Quantization converts model weights from 32-bit floating point to 8-bit integers, reducing model size by approximately 4x and significantly increasing inference speed.
Different quantization strategies offer different speed-accuracy tradeoffs. Dynamic range quantization is the easiest to apply and works well in most cases. Full integer quantization provides the highest speed but requires calibration data.
# Save in Keras format
model.save('final_model.keras')
# Save as TensorFlow SavedModel (for serving)
model.save('saved_model/my_model')
# Convert to TFLite — dynamic range quantization
converter = tf.lite.TFLiteConverter.from_saved_model('saved_model/my_model')
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()
with open('model.tflite', 'wb') as f:
f.write(tflite_model)
# Full integer quantization
def representative_dataset():
"""Generate representative dataset for calibration."""
for i in range(100):
data = np.random.rand(1, 224, 224, 3).astype(np.float32)
yield [data]
converter_int8 = tf.lite.TFLiteConverter.from_saved_model('saved_model/my_model')
converter_int8.optimizations = [tf.lite.Optimize.DEFAULT]
converter_int8.representative_dataset = representative_dataset
converter_int8.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter_int8.inference_input_type = tf.uint8
converter_int8.inference_output_type = tf.uint8
tflite_int8_model = converter_int8.convert()
with open('model_int8.tflite', 'wb') as f:
f.write(tflite_int8_model)
print(f"Original model: {os.path.getsize('model.tflite') / 1024:.1f} KB")
print(f"INT8 model: {os.path.getsize('model_int8.tflite') / 1024:.1f} KB")Deployment Considerations
Taking a model to production is as important as training it. The deployment strategy varies depending on the target platform:
- Web server: Serve via TensorFlow Serving as a REST/gRPC API. Provides high throughput and automatic model versioning.
- Mobile app: Run on Android/iOS with TFLite. Model quantization is essential for size and speed optimization.
- Browser: Client-side inference with TensorFlow.js. No need to send user data to the server.
- Embedded systems: Run on microcontrollers with TFLite Micro. Memory and processing power constraints must be considered.
For every deployment scenario, model versioning, A/B testing, performance monitoring, and rollback mechanisms should be planned. MLOps tools like MLflow, Kubeflow, and Vertex AI can automate these processes.
Summary
Model training with TensorFlow is a combination of efficient data pipelines, proper model architecture, smart callback usage, and systematic hyperparameter search. Custom training loops with GradientTape give you full control over the training process, while mixed precision training boosts GPU efficiency. With transfer learning and a two-phase training strategy, you can achieve high performance even with limited data. TFLite quantization lets you bring your models to mobile devices, reaching a broader audience. Monitoring training with TensorBoard and preserving the best weights with model checkpoints are essential practices for professional ML projects.
Related Articles
What is Machine Learning? Fundamentals & Your First Model with Python
Introduction to machine learning. Supervised, unsupervised, and reinforcement learning concepts. First model with Python, NumPy, and Scikit-learn.
Deep Learning with CNN: Convolutional Neural Network Guide
Convolutional Neural Network (CNN) architecture and deep learning. Convolution, pooling, fully connected layers, and image classification.
Have a Flutter Project?
I build high-performance Flutter applications for iOS, Android, and web.
Get in Touch