Chapter 4: PyTorch and TensorFlow Practice

Learning Objectives

By reading this chapter, you will be able to:

• ✅ Understand PyTorch tensor operations and how automatic differentiation works
• ✅ Build full-fledged neural networks using nn.Module
• ✅ Perform efficient batch processing with DataLoader
• ✅ Compare TensorFlow/Keras implementation patterns with PyTorch
• ✅ Apply practical techniques such as model saving/loading and GPU utilization

4.1 Introduction to PyTorch

What is PyTorch

PyTorch is an open-source machine learning framework developed by Meta (formerly Facebook). It is widely used by researchers and practitioners, and has the following characteristics:

• Define-by-Run: Dynamic computation graphs (the graph is built as the code executes)
• Pythonic: Intuitive, NumPy-like API
• Automatic Differentiation: Automatic gradient computation via Autograd
• GPU Support: Easy acceleration with CUDA

"PyTorch is a framework that can be used seamlessly from research to production. It combines flexibility with high performance."

Tensor Operation Basics

The basic data structure in PyTorch is the tensor (Tensor). It is similar to NumPy's ndarray, but can be computed on the GPU.

import torch
import numpy as np

# Creating tensors
print("=== Creating Tensors ===")

# Create from a list
x = torch.tensor([1.0, 2.0, 3.0])
print(f"x = {x}")
print(f"x.shape = {x.shape}")  # torch.Size([3])

# Initialize with zeros
zeros = torch.zeros(2, 3)
print(f"\nzeros:\n{zeros}")

# Random initialization from a normal distribution
randn = torch.randn(2, 3)
print(f"\nrandn:\n{randn}")

# Convert from a NumPy array
np_array = np.array([[1, 2], [3, 4]])
torch_tensor = torch.from_numpy(np_array)
print(f"\nConverted from NumPy:\n{torch_tensor}")

# Tensor operations
a = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
b = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

print(f"\n=== Tensor Operations ===")
print(f"a + b:\n{a + b}")
print(f"a * b (element-wise):\n{a * b}")
print(f"a @ b (matrix product):\n{a @ b}")
print(f"a.T (transpose):\n{a.T}")

Output:

=== Creating Tensors ===
x = tensor([1., 2., 3.])
x.shape = torch.Size([3])

zeros:
tensor([[0., 0., 0.],
        [0., 0., 0.]])

randn:
tensor([[ 0.3367, -1.2312,  0.5414],
        [-0.8485,  1.1234, -0.3421]])

Converted from NumPy:
tensor([[1, 2],
        [3, 4]])

=== Tensor Operations ===
a + b:
tensor([[ 6.,  8.],
        [10., 12.]])
a * b (element-wise):
tensor([[ 5., 12.],
        [21., 32.]])
a @ b (matrix product):
tensor([[19., 22.],
        [43., 50.]])
a.T (transpose):
tensor([[1., 3.],
        [2., 4.]])

GPU-Enabled Code

import torch

# Check whether a GPU is available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device in use: {device}")

# Transfer a tensor to the GPU
x = torch.randn(1000, 1000)
x_gpu = x.to(device)  # Transfer to GPU
print(f"x_gpu.device: {x_gpu.device}")

# Compute on the GPU
y_gpu = x_gpu @ x_gpu.T
print(f"Shape of GPU computation result: {y_gpu.shape}")

# Move back to the CPU
y_cpu = y_gpu.cpu()
print(f"Transferred to CPU: {y_cpu.device}")

# Create a tensor with an explicit device
z = torch.randn(10, 10, device=device)
print(f"Created with explicit device: {z.device}")

Autograd (Automatic Differentiation)

PyTorch's Autograd tracks tensor operations and automatically computes gradients.

import torch

# Set requires_grad=True on tensors whose gradients you want to compute
x = torch.tensor([2.0], requires_grad=True)
print(f"x = {x}")

# Forward pass: y = x^2 + 3x + 1
y = x**2 + 3*x + 1
print(f"y = {y}")

# Backward pass: compute dy/dx
y.backward()  # Run automatic differentiation

# Retrieve the gradient
print(f"dy/dx = {x.grad}")  # dy/dx = 2x + 3 = 2*2 + 3 = 7

# A more complex example
print("\n=== Complex Computation Graph ===")
x = torch.randn(3, requires_grad=True)
print(f"x = {x}")

# Combine multiple operations
y = x * 2
z = y ** 2
out = z.mean()  # Scalar value

print(f"out = {out}")

# Gradient computation
out.backward()
print(f"x.grad = {x.grad}")

# Verify the gradient: d(mean((x*2)^2))/dx = d(mean(4x^2))/dx = 8x/3
print(f"Expected value (8x/3): {8 * x.data / 3}")

Output:

x = tensor([2.], requires_grad=True)
y = tensor([11.], grad_fn=<AddBackward0>)
dy/dx = tensor([7.])

=== Complex Computation Graph ===
x = tensor([-0.5234,  1.2156, -0.8945], requires_grad=True)
out = tensor(2.0394, grad_fn=<MeanBackward0>)
x.grad = tensor([-1.3957,  3.2416, -2.3853])
Expected value (8x/3): tensor([-1.3957,  3.2416, -2.3853])

Defining Models with nn.Module

nn.Module is the base class for building neural networks in PyTorch.

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleNet(nn.Module):
    """Simple 3-layer neural network"""

    def __init__(self, input_size, hidden_size, output_size):
        """
        Args:
            input_size: Size of the input layer
            hidden_size: Size of the hidden layer
            output_size: Size of the output layer
        """
        super(SimpleNet, self).__init__()

        # Layer definitions
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        """Forward pass"""
        x = F.relu(self.fc1(x))  # 1st layer + ReLU
        x = F.relu(self.fc2(x))  # 2nd layer + ReLU
        x = self.fc3(x)           # Output layer (no activation)
        return x

# Instantiate the model
model = SimpleNet(input_size=784, hidden_size=128, output_size=10)
print(model)

# Check the parameters
print(f"\nTotal parameters: {sum(p.numel() for p in model.parameters()):,}")

# Run inference
x = torch.randn(32, 784)  # Batch size 32
output = model(x)
print(f"\nInput shape: {x.shape}")
print(f"Output shape: {output.shape}")  # (32, 10)

Output:

SimpleNet(
  (fc1): Linear(in_features=784, out_features=128, bias=True)
  (fc2): Linear(in_features=128, out_features=128, bias=True)
  (fc3): Linear(in_features=128, out_features=10, bias=True)
)

Total parameters: 118,026

Input shape: torch.Size([32, 784])
Output shape: torch.Size([32, 10])

Implementing the Training Loop

import torch
import torch.nn as nn
import torch.optim as optim

# Create dummy data
X_train = torch.randn(1000, 784)
y_train = torch.randint(0, 10, (1000,))  # Labels 0-9

# Prepare the model, loss function, and optimizer
model = SimpleNet(784, 128, 10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
num_epochs = 10
batch_size = 32

print("=== Training Started ===")
for epoch in range(num_epochs):
    # Loss per epoch
    epoch_loss = 0.0
    num_batches = 0

    # Split into mini-batches
    for i in range(0, len(X_train), batch_size):
        # Get the batch data
        batch_X = X_train[i:i+batch_size]
        batch_y = y_train[i:i+batch_size]

        # Reset the gradients to zero (important!)
        optimizer.zero_grad()

        # Forward pass
        outputs = model(batch_X)

        # Compute the loss
        loss = criterion(outputs, batch_y)

        # Backward pass
        loss.backward()

        # Update parameters
        optimizer.step()

        epoch_loss += loss.item()
        num_batches += 1

    # Average loss per epoch
    avg_loss = epoch_loss / num_batches
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}")

print("\nTraining complete!")

Output:

=== Training Started ===
Epoch [1/10], Loss: 2.3124
Epoch [2/10], Loss: 2.2845
Epoch [3/10], Loss: 2.2567
Epoch [4/10], Loss: 2.2301
Epoch [5/10], Loss: 2.2048
Epoch [6/10], Loss: 2.1807
Epoch [7/10], Loss: 2.1577
Epoch [8/10], Loss: 2.1357
Epoch [9/10], Loss: 2.1146
Epoch [10/10], Loss: 2.0944

Training complete!

4.2 Practical Model Building

Creating Custom Layers

You can create layers with custom behavior by inheriting from nn.Module.

import torch
import torch.nn as nn
import math

class GaussianNoise(nn.Module):
    """Layer that adds Gaussian noise during training"""

    def __init__(self, sigma=0.1):
        """
        Args:
            sigma: Standard deviation of the noise
        """
        super(GaussianNoise, self).__init__()
        self.sigma = sigma

    def forward(self, x):
        """
        Add noise only during training
        """
        if self.training:  # Only in training mode
            noise = torch.randn_like(x) * self.sigma
            return x + noise
        return x

class ResidualBlock(nn.Module):
    """Block with a residual connection"""

    def __init__(self, size):
        super(ResidualBlock, self).__init__()
        self.fc1 = nn.Linear(size, size)
        self.fc2 = nn.Linear(size, size)
        self.bn1 = nn.BatchNorm1d(size)
        self.bn2 = nn.BatchNorm1d(size)

    def forward(self, x):
        """
        y = F(x) + x  (residual connection)
        """
        residual = x  # Shortcut connection

        out = F.relu(self.bn1(self.fc1(x)))
        out = self.bn2(self.fc2(out))

        out = out + residual  # Add the residual
        out = F.relu(out)

        return out

# Model incorporating the custom layers
class CustomNet(nn.Module):
    def __init__(self):
        super(CustomNet, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.noise = GaussianNoise(sigma=0.05)
        self.res1 = ResidualBlock(256)
        self.res2 = ResidualBlock(256)
        self.fc_out = nn.Linear(256, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.noise(x)  # Add noise (regularization)
        x = self.res1(x)   # Residual block 1
        x = self.res2(x)   # Residual block 2
        x = self.fc_out(x)
        return x

# Test
model = CustomNet()
x = torch.randn(16, 784)

# Training mode
model.train()
out_train = model(x)
print(f"Training mode output: {out_train.shape}")

# Evaluation mode
model.eval()
out_eval = model(x)
print(f"Evaluation mode output: {out_eval.shape}")

Loss Functions and Optimizers

PyTorch provides a wide variety of loss functions and optimizers.

import torch
import torch.nn as nn
import torch.optim as optim

# Prepare the model
model = SimpleNet(784, 128, 10)

print("=== Choosing a Loss Function ===")

# Loss functions for classification problems
criterion_ce = nn.CrossEntropyLoss()  # Multi-class classification
criterion_bce = nn.BCEWithLogitsLoss()  # Binary classification
criterion_nll = nn.NLLLoss()  # Negative log-likelihood

# Loss functions for regression problems
criterion_mse = nn.MSELoss()  # Mean squared error
criterion_mae = nn.L1Loss()   # Mean absolute error
criterion_huber = nn.SmoothL1Loss()  # Huber loss

print("Classification: CrossEntropyLoss, BCEWithLogitsLoss, NLLLoss")
print("Regression: MSELoss, L1Loss, SmoothL1Loss")

print("\n=== Choosing an Optimizer ===")

# SGD (stochastic gradient descent)
optimizer_sgd = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# Adam (adaptive learning rate)
optimizer_adam = optim.Adam(model.parameters(), lr=0.001, betas=(0.9, 0.999))

# AdamW (Adam with improved weight decay)
optimizer_adamw = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)

# RMSprop
optimizer_rmsprop = optim.RMSprop(model.parameters(), lr=0.01, alpha=0.99)

print("SGD: classic, stabilized with momentum")
print("Adam: adaptive learning rate, good by default")
print("AdamW: improved version of Adam, better regularization")
print("RMSprop: effective for RNNs")

# Example of using an optimizer
print("\n=== Example of a Training Step ===")
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Dummy data
x = torch.randn(32, 784)
y = torch.randint(0, 10, (32,))

# One training step
optimizer.zero_grad()      # Reset gradients
output = model(x)          # Forward pass
loss = criterion(output, y)  # Compute loss
loss.backward()            # Backward pass
optimizer.step()           # Update parameters

print(f"Loss: {loss.item():.4f}")

DataLoader and Batch Processing

DataLoader provides automatic batching, shuffling, and parallel data loading.

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

# Custom dataset class
class CustomDataset(Dataset):
    """Example implementation of a custom dataset"""

    def __init__(self, data, labels, transform=None):
        """
        Args:
            data: Input data (NumPy array or tensor)
            labels: Label data
            transform: Data transformation function (optional)
        """
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        """Return the size of the dataset"""
        return len(self.data)

    def __getitem__(self, idx):
        """Return the data corresponding to the index"""
        x = self.data[idx]
        y = self.labels[idx]

        # Apply the transformation
        if self.transform:
            x = self.transform(x)

        return x, y

# Create dummy data
np.random.seed(42)
data = np.random.randn(1000, 784).astype(np.float32)
labels = np.random.randint(0, 10, size=1000)

# Create the dataset
dataset = CustomDataset(data, labels)

# Create the DataLoader
train_loader = DataLoader(
    dataset,
    batch_size=32,      # Batch size
    shuffle=True,       # Shuffle every epoch
    num_workers=0,      # Number of parallel data-loading workers (0 = main process only)
    drop_last=True      # Drop the last incomplete batch
)

print(f"Dataset size: {len(dataset)}")
print(f"Number of batches: {len(train_loader)}")

# Example of using the DataLoader
print("\n=== Batch Processing with DataLoader ===")
for batch_idx, (data_batch, labels_batch) in enumerate(train_loader):
    if batch_idx >= 3:  # Show only the first 3 batches
        break
    print(f"Batch {batch_idx+1}: data shape={data_batch.shape}, labels shape={labels_batch.shape}")

# Example of use in a training loop
print("\n=== Training Loop Example ===")
model = SimpleNet(784, 128, 10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(3):
    model.train()
    total_loss = 0

    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    avg_loss = total_loss / len(train_loader)
    print(f"Epoch {epoch+1}: Average loss = {avg_loss:.4f}")

Visualizing Training

import matplotlib.pyplot as plt
from IPython.display import clear_output

class TrainingMonitor:
    """Class that visualizes the training process"""

    def __init__(self):
        self.train_losses = []
        self.val_losses = []
        self.train_accs = []
        self.val_accs = []

    def update(self, train_loss, val_loss, train_acc, val_acc):
        """Append the metrics"""
        self.train_losses.append(train_loss)
        self.val_losses.append(val_loss)
        self.train_accs.append(train_acc)
        self.val_accs.append(val_acc)

    def plot(self):
        """Real-time plotting"""
        clear_output(wait=True)

        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))

        # Plot the loss
        ax1.plot(self.train_losses, label='Train Loss', marker='o')
        ax1.plot(self.val_losses, label='Val Loss', marker='s')
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss')
        ax1.set_title('Training and Validation Loss')
        ax1.legend()
        ax1.grid(True, alpha=0.3)

        # Plot the accuracy
        ax2.plot(self.train_accs, label='Train Acc', marker='o')
        ax2.plot(self.val_accs, label='Val Acc', marker='s')
        ax2.set_xlabel('Epoch')
        ax2.set_ylabel('Accuracy (%)')
        ax2.set_title('Training and Validation Accuracy')
        ax2.legend()
        ax2.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.show()

# Usage example
monitor = TrainingMonitor()

# Simulate training with dummy data
for epoch in range(10):
    # Loss gradually decreases
    train_loss = 2.0 - epoch * 0.15 + np.random.randn() * 0.05
    val_loss = 2.1 - epoch * 0.13 + np.random.randn() * 0.05

    # Accuracy gradually improves
    train_acc = 20 + epoch * 7 + np.random.randn() * 2
    val_acc = 18 + epoch * 6.5 + np.random.randn() * 2

    monitor.update(train_loss, val_loss, train_acc, val_acc)
    monitor.plot()

print("Training curve visualization complete")

4.3 Introduction to TensorFlow/Keras

Overview of TensorFlow and Keras

TensorFlow is a deep learning framework developed by Google, and Keras is integrated into it as a high-level API.

• Define-and-Run: Static computation graphs (advantageous for acceleration)
• Eager Execution: TensorFlow 2.x and later also support dynamic execution
• Production Ready: Easy production deployment via TensorFlow Serving, TF Lite, and TF.js
• Keras API: Simple and intuitive API

Sequential API and Functional API

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import numpy as np

print(f"TensorFlow version: {tf.__version__}")

# ===== Sequential API =====
print("\n=== Sequential API ===")
# Simple model that stacks layers in order

model_seq = keras.Sequential([
    layers.Dense(128, activation='relu', input_shape=(784,)),
    layers.Dropout(0.2),
    layers.Dense(128, activation='relu'),
    layers.Dropout(0.2),
    layers.Dense(10, activation='softmax')
])

model_seq.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

print(model_seq.summary())

# ===== Functional API =====
print("\n=== Functional API ===")
# More complex structures possible (branches, multiple inputs/outputs)

# Define the input
inputs = keras.Input(shape=(784,))

# Define the hidden layers
x = layers.Dense(128, activation='relu')(inputs)
x = layers.Dropout(0.2)(x)
x = layers.Dense(128, activation='relu')(x)
x = layers.Dropout(0.2)(x)

# Output layer
outputs = layers.Dense(10, activation='softmax')(x)

# Build the model
model_func = keras.Model(inputs=inputs, outputs=outputs)

model_func.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

print(model_func.summary())

# Train with dummy data
X_train = np.random.randn(1000, 784).astype(np.float32)
y_train = np.random.randint(0, 10, size=1000)

print("\n=== Running Training ===")
history = model_seq.fit(
    X_train, y_train,
    batch_size=32,
    epochs=5,
    validation_split=0.2,
    verbose=1
)

Creating Custom Models

import tensorflow as tf
from tensorflow import keras

class CustomModel(keras.Model):
    """Custom model class (equivalent to PyTorch's nn.Module)"""

    def __init__(self, num_classes=10):
        super(CustomModel, self).__init__()
        self.dense1 = layers.Dense(128, activation='relu')
        self.dropout1 = layers.Dropout(0.2)
        self.dense2 = layers.Dense(128, activation='relu')
        self.dropout2 = layers.Dropout(0.2)
        self.dense3 = layers.Dense(num_classes, activation='softmax')

    def call(self, inputs, training=False):
        """Forward pass (equivalent to PyTorch's forward)"""
        x = self.dense1(inputs)
        x = self.dropout1(x, training=training)  # Dropout only during training
        x = self.dense2(x)
        x = self.dropout2(x, training=training)
        return self.dense3(x)

# Instantiate the custom model
custom_model = CustomModel(num_classes=10)

# Compile
custom_model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# Train with dummy data
X_train = np.random.randn(1000, 784).astype(np.float32)
y_train = np.random.randint(0, 10, size=1000)

history = custom_model.fit(
    X_train, y_train,
    batch_size=32,
    epochs=5,
    validation_split=0.2,
    verbose=1
)

print("\nCustom model training complete")

Using Callbacks

Callbacks are routines executed at specific points during training.

from tensorflow.keras.callbacks import (
    EarlyStopping,
    ModelCheckpoint,
    ReduceLROnPlateau,
    TensorBoard
)
import datetime

# Early Stopping: stop training when the validation loss stops improving
early_stop = EarlyStopping(
    monitor='val_loss',     # Metric to monitor
    patience=5,              # Number of epochs without improvement
    restore_best_weights=True  # Restore the best weights
)

# Model Checkpoint: save the best model
checkpoint = ModelCheckpoint(
    filepath='best_model.h5',
    monitor='val_accuracy',
    save_best_only=True,
    verbose=1
)

# Learning Rate Scheduler: adjust the learning rate dynamically
reduce_lr = ReduceLROnPlateau(
    monitor='val_loss',
    factor=0.5,              # Halve the learning rate
    patience=3,
    min_lr=1e-6,
    verbose=1
)

# TensorBoard: visualize the training process
log_dir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=1)

# Training with callbacks
model = keras.Sequential([
    layers.Dense(128, activation='relu', input_shape=(784,)),
    layers.Dropout(0.2),
    layers.Dense(64, activation='relu'),
    layers.Dense(10, activation='softmax')
])

model.compile(
    optimizer='adam',
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# Dummy data
X_train = np.random.randn(1000, 784).astype(np.float32)
y_train = np.random.randint(0, 10, size=1000)
X_val = np.random.randn(200, 784).astype(np.float32)
y_val = np.random.randint(0, 10, size=200)

# Train with the specified callbacks
history = model.fit(
    X_train, y_train,
    batch_size=32,
    epochs=20,
    validation_data=(X_val, y_val),
    callbacks=[early_stop, checkpoint, reduce_lr],
    verbose=1
)

print("\nTraining with callbacks complete")
print(f"Final training epoch: {len(history.history['loss'])}")

Comparison with PyTorch

How to choose:
- PyTorch: research, priority on flexibility, fine-grained control needed
- TensorFlow/Keras: production deployment, priority on speed, simple API

4.4 Practical Techniques

GPU Usage and Memory Management

import torch
import torch.nn as nn

# === GPU usage basics ===
print("=== GPU Information ===")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"Number of GPUs: {torch.cuda.device_count()}")
    print(f"GPU name: {torch.cuda.get_device_name(0)}")
    print(f"Memory allocated: {torch.cuda.memory_allocated(0) / 1024**2:.2f} MB")
    print(f"Memory reserved: {torch.cuda.memory_reserved(0) / 1024**2:.2f} MB")

# Select the device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"\nDevice in use: {device}")

# Transfer the model and data to the GPU
model = SimpleNet(784, 128, 10).to(device)
x = torch.randn(32, 784, device=device)  # Create directly on the GPU

# Inference
with torch.no_grad():  # Disable gradient computation (saves memory)
    output = model(x)
print(f"Output device: {output.device}")

# === Memory management best practices ===
print("\n=== Memory Management ===")

# 1. Delete unneeded tensors
x = torch.randn(1000, 1000, device=device)
del x
torch.cuda.empty_cache()  # Clear the cache

# 2. Gradient Checkpointing (trade-off between memory and computation)
from torch.utils.checkpoint import checkpoint

class MemoryEfficientNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer1 = nn.Linear(784, 1024)
        self.layer2 = nn.Linear(1024, 1024)
        self.layer3 = nn.Linear(1024, 10)

    def forward(self, x):
        # Using checkpoint reduces memory usage (at the cost of slower computation)
        x = checkpoint(lambda x: F.relu(self.layer1(x)), x)
        x = checkpoint(lambda x: F.relu(self.layer2(x)), x)
        x = self.layer3(x)
        return x

# 3. Mixed Precision Training (memory reduction + speedup)
from torch.cuda.amp import autocast, GradScaler

model = SimpleNet(784, 128, 10).to(device)
optimizer = torch.optim.Adam(model.parameters())
scaler = GradScaler()  # Automatic scaling

for epoch in range(3):
    for data, target in [(torch.randn(32, 784, device=device),
                          torch.randint(0, 10, (32,), device=device))]:
        optimizer.zero_grad()

        # Mixed Precision
        with autocast():
            output = model(data)
            loss = F.cross_entropy(output, target)

        # Backward pass with scaling
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

    print(f"Epoch {epoch+1} complete")

print("\nGPU-optimized training complete")

Saving and Loading Models

import torch
import torch.nn as nn

# Prepare the model
model = SimpleNet(784, 128, 10)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# ===== Saving and loading with PyTorch =====
print("=== PyTorch: Saving Models ===")

# Method 1: Save only the state_dict (recommended)
torch.save(model.state_dict(), 'model_weights.pth')
print("Weights saved: model_weights.pth")

# Method 2: Save the entire model
torch.save(model, 'model_full.pth')
print("Entire model saved: model_full.pth")

# Method 3: Checkpoint (mid-training state)
checkpoint = {
    'epoch': 10,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': 0.5,
}
torch.save(checkpoint, 'checkpoint.pth')
print("Checkpoint saved: checkpoint.pth")

print("\n=== PyTorch: Loading Models ===")

# Method 1: Load the state_dict
model_new = SimpleNet(784, 128, 10)
model_new.load_state_dict(torch.load('model_weights.pth'))
model_new.eval()  # Set to evaluation mode
print("Weights loaded")

# Method 2: Load the entire model
model_loaded = torch.load('model_full.pth')
model_loaded.eval()
print("Entire model loaded")

# Method 3: Resume from a checkpoint
checkpoint = torch.load('checkpoint.pth')
model_new.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']
print(f"Checkpoint loaded: Epoch={epoch}, Loss={loss}")

# ===== Saving and loading with TensorFlow/Keras =====
print("\n=== TensorFlow/Keras: Saving Models ===")

# Prepare a Keras model
keras_model = keras.Sequential([
    layers.Dense(128, activation='relu', input_shape=(784,)),
    layers.Dense(10, activation='softmax')
])
keras_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')

# Method 1: SavedFormat (recommended)
keras_model.save('saved_model/')
print("SavedFormat saved: saved_model/")

# Method 2: HDF5 format
keras_model.save('model.h5')
print("HDF5 saved: model.h5")

# Method 3: Weights only
keras_model.save_weights('weights.h5')
print("Weights only saved: weights.h5")

print("\n=== TensorFlow/Keras: Loading Models ===")

# Load from SavedFormat
loaded_model = keras.models.load_model('saved_model/')
print("SavedFormat loaded")

# Load from HDF5
loaded_model_h5 = keras.models.load_model('model.h5')
print("HDF5 loaded")

# Load weights only
new_model = keras.Sequential([
    layers.Dense(128, activation='relu', input_shape=(784,)),
    layers.Dense(10, activation='softmax')
])
new_model.load_weights('weights.h5')
print("Weights only loaded")

Early Stopping and Checkpointing

import torch
import numpy as np

class EarlyStopping:
    """Class implementing Early Stopping (for PyTorch)"""

    def __init__(self, patience=7, min_delta=0, restore_best_weights=True):
        """
        Args:
            patience: Number of epochs without improvement
            min_delta: Minimum change to qualify as an improvement
            restore_best_weights: Whether to restore the best weights
        """
        self.patience = patience
        self.min_delta = min_delta
        self.restore_best_weights = restore_best_weights

        self.counter = 0
        self.best_loss = None
        self.early_stop = False
        self.best_weights = None

    def __call__(self, val_loss, model):
        """
        Args:
            val_loss: Validation loss
            model: The model

        Returns:
            Whether to stop early
        """
        if self.best_loss is None:
            self.best_loss = val_loss
            self.save_checkpoint(model)
        elif val_loss > self.best_loss - self.min_delta:
            self.counter += 1
            print(f'EarlyStopping counter: {self.counter}/{self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
                if self.restore_best_weights:
                    print('Restoring the best weights')
                    model.load_state_dict(self.best_weights)
        else:
            self.best_loss = val_loss
            self.save_checkpoint(model)
            self.counter = 0

        return self.early_stop

    def save_checkpoint(self, model):
        """Save the best model"""
        if self.restore_best_weights:
            self.best_weights = model.state_dict().copy()

# Usage example
print("=== Early Stopping Implementation Example ===")
model = SimpleNet(784, 128, 10)
optimizer = torch.optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
early_stopping = EarlyStopping(patience=5, restore_best_weights=True)

# Dummy data
X_train = torch.randn(800, 784)
y_train = torch.randint(0, 10, (800,))
X_val = torch.randn(200, 784)
y_val = torch.randint(0, 10, (200,))

for epoch in range(50):
    # Training
    model.train()
    optimizer.zero_grad()
    output = model(X_train)
    loss = criterion(output, y_train)
    loss.backward()
    optimizer.step()

    # Validation
    model.eval()
    with torch.no_grad():
        val_output = model(X_val)
        val_loss = criterion(val_output, y_val)

    print(f'Epoch {epoch+1}: Train Loss={loss.item():.4f}, Val Loss={val_loss.item():.4f}')

    # Early Stopping check
    if early_stopping(val_loss.item(), model):
        print(f'Early stopping at epoch {epoch+1}')
        break

print("\nTraining complete")

Learning Rate Scheduling

import torch
import torch.optim as optim
from torch.optim.lr_scheduler import (
    StepLR,
    ExponentialLR,
    CosineAnnealingLR,
    ReduceLROnPlateau
)
import matplotlib.pyplot as plt

# Model and optimizer
model = SimpleNet(784, 128, 10)
optimizer = optim.Adam(model.parameters(), lr=0.1)

print("=== Types of Learning Rate Schedulers ===\n")

# 1. StepLR: decay the learning rate every fixed number of epochs
scheduler1 = StepLR(optimizer, step_size=10, gamma=0.5)
print("StepLR: halve the learning rate every 10 epochs")

# 2. ExponentialLR: exponential decay
scheduler2 = ExponentialLR(optimizer, gamma=0.95)
print("ExponentialLR: 5% decay every epoch")

# 3. CosineAnnealingLR: decay following a cosine function
scheduler3 = CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)
print("CosineAnnealingLR: smooth decay following a cosine function")

# 4. ReduceLROnPlateau: decay when the loss stops improving
scheduler4 = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)
print("ReduceLROnPlateau: decay when the loss stops improving")

# Visualize the learning rate trajectory
print("\n=== Learning Rate Trajectory ===")

schedulers = {
    'StepLR': StepLR(optim.Adam(model.parameters(), lr=0.1), step_size=10, gamma=0.5),
    'ExponentialLR': ExponentialLR(optim.Adam(model.parameters(), lr=0.1), gamma=0.95),
    'CosineAnnealingLR': CosineAnnealingLR(optim.Adam(model.parameters(), lr=0.1), T_max=50)
}

fig, ax = plt.subplots(figsize=(10, 6))

for name, scheduler in schedulers.items():
    lrs = []
    optimizer = scheduler.optimizer
    for epoch in range(50):
        lrs.append(optimizer.param_groups[0]['lr'])
        scheduler.step()

    ax.plot(lrs, label=name, linewidth=2)

ax.set_xlabel('Epoch', fontsize=12)
ax.set_ylabel('Learning Rate', fontsize=12)
ax.set_title('Comparison of Learning Rate Schedulers', fontsize=14, fontweight='bold')
ax.legend(fontsize=11)
ax.grid(True, alpha=0.3)
ax.set_yscale('log')
plt.tight_layout()
plt.show()

# Practical usage example
print("\n=== Scheduler Usage Example ===")
model = SimpleNet(784, 128, 10)
optimizer = optim.Adam(model.parameters(), lr=0.01)
scheduler = CosineAnnealingLR(optimizer, T_max=10)

for epoch in range(10):
    # Training step (omitted)
    current_lr = optimizer.param_groups[0]['lr']
    print(f"Epoch {epoch+1}: Learning Rate = {current_lr:.6f}")

    # Update the scheduler at the end of the epoch
    scheduler.step()

print("\nLearning rate scheduling complete")

4.5 Chapter Summary

What We Learned

1. PyTorch Basics

   • Tensor operations and GPU utilization
   • Automatic differentiation with Autograd
   • Model definition with nn.Module
   • Training loop implementation patterns

2. Practical Model Building

   • Custom layers (residual connections, noise layers)
   • Choosing loss functions and optimizers
   • Efficient batch processing with DataLoader
   • Visualizing the training process

3. TensorFlow/Keras

   • The three APIs: Sequential/Functional/Custom Model
   • Controlling training with callbacks
   • Differences from PyTorch and when to use each

4. Practical Techniques

   • GPU memory management and Mixed Precision
   • Best practices for saving and loading models
   • Early Stopping and Checkpointing
   • Learning rate scheduling

Framework Selection Guide

Best Practices

1. Device management: Write code that works regardless of GPU availability
2. Reproducibility: Fix the random seed (torch.manual_seed(42))
3. Memory efficiency: Save memory during inference with torch.no_grad()
4. Training monitoring: Visualize with TensorBoard or WandB
5. Checkpointing: Save regularly to guard against training interruptions

To the Next Chapter

In Chapter 5, we will learn about Convolutional Neural Networks (CNNs):

• Convolutional layers for image recognition
• Pooling layers and feature extraction
• Representative CNN architectures (LeNet, VGG, ResNet)
• Implementation on real data (CIFAR-10, ImageNet)

Exercises

Problem 1 (Difficulty: easy)

Determine whether each of the following statements is true or false.

1. PyTorch adopts the Define-by-Run approach
2. In TensorFlow 2.x, Eager Execution is enabled by default
3. The forward method of nn.Module must be called manually
4. DataLoader's shuffle=True shuffles the data order every epoch

Problem 2 (Difficulty: medium)

Fix the bugs in the following PyTorch code.

import torch
import torch.nn as nn

model = nn.Linear(10, 5)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for epoch in range(10):
    x = torch.randn(32, 10)
    y = torch.randn(32, 5)

    output = model(x)
    loss = nn.MSELoss(output, y)
    loss.backward()
    optimizer.step()

import torch
import torch.nn as nn

model = nn.Linear(10, 5)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()  # Instantiate the loss function

for epoch in range(10):
    x = torch.randn(32, 10)
    y = torch.randn(32, 5)

    optimizer.zero_grad()  # Reset the gradients to zero (added)
    output = model(x)
    loss = criterion(output, y)  # Fixed
    loss.backward()
    optimizer.step()

Problem 3 (Difficulty: medium)

Implement a custom dataset class that satisfies the following requirements:

• Accepts NumPy arrays for the data and labels
• Normalizes the data (mean 0, standard deviation 1)
• Usable with DataLoader

import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np

class NormalizedDataset(Dataset):
    """Custom dataset that applies normalization"""

    def __init__(self, data, labels):
        """
        Args:
            data: NumPy array (N, features)
            labels: NumPy array (N,)
        """
        # Normalize the data
        self.mean = data.mean(axis=0)
        self.std = data.std(axis=0) + 1e-8  # Avoid division by zero
        self.data = (data - self.mean) / self.std
        self.labels = labels

        # Convert to PyTorch tensors
        self.data = torch.FloatTensor(self.data)
        self.labels = torch.LongTensor(self.labels)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]

# Test
np.random.seed(42)
data = np.random.randn(100, 10) * 5 + 10  # Mean 10, standard deviation 5
labels = np.random.randint(0, 3, size=100)

dataset = NormalizedDataset(data, labels)
loader = DataLoader(dataset, batch_size=16, shuffle=True)

# Verify
sample_data, sample_label = dataset[0]
print(f"Normalized data (first sample): {sample_data}")
print(f"Mean of data (after normalization): {dataset.data.mean():.4f}")
print(f"Standard deviation of data (after normalization): {dataset.data.std():.4f}")

Problem 4 (Difficulty: hard)

Implement the same model architecture in both PyTorch and TensorFlow/Keras:

• Input: 28x28 image (784 dimensions)
• Hidden layer 1: 256 units, ReLU, Dropout(0.3)
• Hidden layer 2: 128 units, ReLU, Dropout(0.3)
• Output layer: 10-class classification

import torch
import torch.nn as nn
import torch.nn.functional as F

class PyTorchModel(nn.Module):
    def __init__(self):
        super(PyTorchModel, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.dropout1 = nn.Dropout(0.3)
        self.fc2 = nn.Linear(256, 128)
        self.dropout2 = nn.Dropout(0.3)
        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = self.dropout1(x)
        x = F.relu(self.fc2(x))
        x = self.dropout2(x)
        x = self.fc3(x)
        return x

# Usage example
pytorch_model = PyTorchModel()
x = torch.randn(32, 784)
output = pytorch_model(x)
print(f"PyTorch output shape: {output.shape}")

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

# Sequential API
keras_model = keras.Sequential([
    layers.Dense(256, activation='relu', input_shape=(784,)),
    layers.Dropout(0.3),
    layers.Dense(128, activation='relu'),
    layers.Dropout(0.3),
    layers.Dense(10)
])

# Usage example
x = tf.random.normal((32, 784))
output = keras_model(x)
print(f"Keras output shape: {output.shape}")

# PyTorch
pytorch_params = sum(p.numel() for p in pytorch_model.parameters())
print(f"PyTorch parameters: {pytorch_params:,}")

# Keras
keras_model.summary()

Problem 5 (Difficulty: hard)

Implement training code that combines Early Stopping and Model Checkpointing. It must satisfy the following conditions:

• Stop when the validation loss does not improve for 5 epochs
• Automatically save the best model
• Restore the best model at the end of training

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

class EarlyStoppingWithCheckpoint:
    """Class that integrates Early Stopping and Checkpointing"""

    def __init__(self, patience=5, checkpoint_path='best_model.pth', verbose=True):
        self.patience = patience
        self.checkpoint_path = checkpoint_path
        self.verbose = verbose

        self.counter = 0
        self.best_loss = None
        self.early_stop = False

    def __call__(self, val_loss, model):
        if self.best_loss is None or val_loss < self.best_loss:
            # Improvement
            if self.verbose:
                if self.best_loss is not None:
                    improvement = self.best_loss - val_loss
                    print(f'Validation loss improved: {self.best_loss:.4f} → {val_loss:.4f} ({improvement:.4f})')
                else:
                    print(f'Initial best model: {val_loss:.4f}')

            self.best_loss = val_loss
            self.counter = 0

            # Save the model
            torch.save(model.state_dict(), self.checkpoint_path)
            if self.verbose:
                print(f'Model saved: {self.checkpoint_path}')
        else:
            # No improvement
            self.counter += 1
            if self.verbose:
                print(f'EarlyStopping counter: {self.counter}/{self.patience}')

            if self.counter >= self.patience:
                self.early_stop = True
                if self.verbose:
                    print('Early Stopping triggered!')

        return self.early_stop

    def load_best_model(self, model):
        """Load the best model"""
        model.load_state_dict(torch.load(self.checkpoint_path))
        if self.verbose:
            print(f'Best model restored: {self.checkpoint_path} (loss={self.best_loss:.4f})')

# Usage example
print("=== Early Stopping + Checkpoint Implementation Example ===\n")

# Dummy data
X_train = torch.randn(800, 784)
y_train = torch.randint(0, 10, (800,))
X_val = torch.randn(200, 784)
y_val = torch.randint(0, 10, (200,))

# Model, loss function, and optimizer
model = SimpleNet(784, 128, 10)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Initialize Early Stopping
early_stopping = EarlyStoppingWithCheckpoint(
    patience=5,
    checkpoint_path='best_model.pth',
    verbose=True
)

# Training loop
for epoch in range(50):
    # Training
    model.train()
    optimizer.zero_grad()
    output = model(X_train)
    train_loss = criterion(output, y_train)
    train_loss.backward()
    optimizer.step()

    # Validation
    model.eval()
    with torch.no_grad():
        val_output = model(X_val)
        val_loss = criterion(val_output, y_val)

    print(f'\nEpoch {epoch+1}: Train Loss={train_loss.item():.4f}, Val Loss={val_loss.item():.4f}')

    # Early Stopping check
    if early_stopping(val_loss.item(), model):
        print(f'\nTraining stopped at epoch {epoch+1}')
        break

# Restore the best model
print("\n" + "="*50)
early_stopping.load_best_model(model)
print("="*50)

References

1. Paszke, A., et al. (2019). "PyTorch: An Imperative Style, High-Performance Deep Learning Library." NeurIPS.
2. Abadi, M., et al. (2016). "TensorFlow: A System for Large-Scale Machine Learning." OSDI.
3. Chollet, F. (2017). Deep Learning with Python. Manning Publications.
4. Stevens, E., Antiga, L., & Viehmann, T. (2020). Deep Learning with PyTorch. Manning Publications.