Deep Learning Models: A collection of various deep learning architectures, models, & their implementations (PyTorch or TensorFlow) in Jupyter Notebooks.

• Author: Chizoba Obasi
• Github Repo: https://github.com/chizkidd/deep-learning-models

!pip install watermark

Collecting watermark
  Downloading watermark-2.5.1-py3-none-any.whl.metadata (1.5 kB)
Requirement already satisfied: ipython>=6.0 in /usr/local/lib/python3.11/dist-packages (from watermark) (7.34.0)
Requirement already satisfied: importlib-metadata>=1.4 in /usr/local/lib/python3.11/dist-packages (from watermark) (8.7.0)
Requirement already satisfied: setuptools in /usr/local/lib/python3.11/dist-packages (from watermark) (75.2.0)
Requirement already satisfied: zipp>=3.20 in /usr/local/lib/python3.11/dist-packages (from importlib-metadata>=1.4->watermark) (3.23.0)
Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.11/dist-packages (from ipython>=6.0->watermark) (0.19.2)
Requirement already satisfied: decorator in /usr/local/lib/python3.11/dist-packages (from ipython>=6.0->watermark) (4.4.2)
Requirement already satisfied: pickleshare in /usr/local/lib/python3.11/dist-packages (from ipython>=6.0->watermark) (0.7.5)
Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.11/dist-packages (from ipython>=6.0->watermark) (5.7.1)
Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.11/dist-packages (from ipython>=6.0->watermark) (3.0.51)
Requirement already satisfied: pygments in /usr/local/lib/python3.11/dist-packages (from ipython>=6.0->watermark) (2.19.2)
Requirement already satisfied: backcall in /usr/local/lib/python3.11/dist-packages (from ipython>=6.0->watermark) (0.2.0)
Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.11/dist-packages (from ipython>=6.0->watermark) (0.1.7)
Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.11/dist-packages (from ipython>=6.0->watermark) (4.9.0)
Requirement already satisfied: parso<0.9.0,>=0.8.4 in /usr/local/lib/python3.11/dist-packages (from jedi>=0.16->ipython>=6.0->watermark) (0.8.4)
Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.11/dist-packages (from pexpect>4.3->ipython>=6.0->watermark) (0.7.0)
Requirement already satisfied: wcwidth in /usr/local/lib/python3.11/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython>=6.0->watermark) (0.2.13)
Downloading watermark-2.5.1-py3-none-any.whl (7.8 kB)
Installing collected packages: watermark
Successfully installed watermark-2.5.1

%load_ext watermark
# %reload_ext watermark
%watermark -a 'Chizoba Obasi' -v -p torch

Author: Chizoba Obasi

Python implementation: CPython
Python version       : 3.11.13
IPython version      : 7.34.0

torch: 2.6.0+cu124

VGG-16 Classifier on CIFAR-10

Network Architecture:

Image Source

ConveNet Configuration:

Network has 5 convolutional stages:

1. conv->conv
2. conv->conv
3. conv->conv->conv->[conv]
4. conv->conv->conv->[conv]
5. conv->conv->conv->[conv]
6. FC->FC->FC

• VGG-16 has 16 layers if you count all the conv and FC layers except the bracketed conv [conv].
• VGG-19 has the extra bracketed conv [conv] in stages 3, 4 and 5 .

Design Rules:

1. All conv are $3\times3$ withstride of $1$ and padding of $1$
2. All maxpool are $2\times2$ with stride of $2$
3. After maxpool, double # of channels

Notes:

1. Two $3\times3$ conv has same receptive field as a single $5\times5$ conv, but has fewer parameters and takes less computation!

Options	1	2
Setup	CONV($5\times5$, $C \rightarrow C$)	2x CONV($3\times3$, $C \rightarrow C$)
Params	$25C^2$	$2 \times 9C^2 = 18C^2$
FLOPs	$25C^2HW$	$18C^2HW$

2. conv layers at each spatial resolution take the same amount of computation!

Config	1	2
Input	$C \times 2H \times 2W$	$2C \times H \times W$
Layer	CONV($3\times3$, $C \rightarrow C$)	CONV($3\times3$, $2C \rightarrow 2C$)
Memory	$4HWC$	$2HWC$
Params	$9C^2$	$2 \times 9C^2 = 36C^2$
FLOPs	$36C^2HW$	$36C^2HW$

REFERENCES:

1. Simonyan, Karen, and Andrew Zisserman. "Very Deep Convolutional Networks for Large-Scale Image Recognition." International Conference on Learning Representations (ICLR), 2015. arXiv:1409.1556.
2. Johnson, J. (2022, February 2). Lecture 8: CNN Architectures [PowerPoint Slides]. University of Michigan EECS 498/598: Deep Learning for Computer Vision. pp 56-65. https://web.eecs.umich.edu/~justincj/slides/eecs498/WI2022/598_WI2022_lecture08.pdf

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from sklearn.model_selection import train_test_split

import matplotlib.image as mpimg
import matplotlib.pyplot as plt
%matplotlib inline

import pandas as pd
import numpy as np
import math
import time
import os
import random

if torch.cuda.is_available():
    torch.backends.cudnn.deterministic = True

!pip install torchinfo

Requirement already satisfied: torchinfo in /usr/local/lib/python3.11/dist-packages (1.8.0)

from torchinfo import summary

def set_all_seeds(seed):
    os.environ["PL_GLOBAL_SEED"] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def set_deterministic():
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
    torch.set_deterministic(True)

# Check if GPU is available and set the device
if torch.cuda.is_available():
    num_gpus = torch.cuda.device_count()
    print("Num GPUs Available: ", num_gpus)
    for i in range(num_gpus):
        gpu_name = torch.cuda.get_device_name(i)
        print(f"GPU {i}: {gpu_name}")
else:
    print("No GPU available.")

Num GPUs Available:  2
GPU 0: Tesla T4
GPU 1: Tesla T4

---

0. Model Hyperparameters

---

##########################
### SETTINGS
##########################

# Hyperparameters
RANDOM_SEED = 123
BATCH_SIZE = 128 #256
LEARNING_RATE = 0.001
NUM_EPOCHS = 20
INPUT_SIZE = 32 # CIFAR-10 is 32x32
NUM_CLASSES = 10

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

cuda

---

1. Data Preparation & Loading

---

# Mean and Std Dev for CIFAR-10 (3 channels)
CIFAR_MEAN = (0.4914, 0.4822, 0.4465)
CIFAR_STD = (0.2471, 0.2435, 0.2616)

# Data Augmentation (Crucial for VGG-16 on CIFAR-10)
train_transforms = transforms.Compose([
    transforms.RandomCrop(32, padding=4),   # Crop after padding
    transforms.RandomHorizontalFlip(),      # Flip
    transforms.ToTensor(),
    transforms.Normalize(CIFAR_MEAN, CIFAR_STD) # Mean Subtraction
])

# Test/Validation Transforms (No Augmentation)
test_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(CIFAR_MEAN, CIFAR_STD)
])

# Load Datasets
train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transforms)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=test_transforms)

# Split training set into training and validation sets
train_size = int(0.9 * len(train_dataset))
val_size = len(train_dataset) - train_size
train_subset, val_subset = torch.utils.data.random_split(train_dataset, [train_size, val_size])

# Create DataLoaders
train_loader = DataLoader(train_subset, batch_size=BATCH_SIZE, shuffle=True)
val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)

100%|██████████| 170M/170M [00:04<00:00, 35.4MB/s] 

---

2. Model Creation

---

# VGG-16 Architecture (List of output channels for each layer)
# M denotes a MaxPool2d layer
VGG16_arch = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M']

class VGG(nn.Module):
    def __init__(self, features, num_classes=10):
        super(VGG, self).__init__()
        
        # Convolutional Block
        self.features = features

        # 2x2 MaxPool (stride=2): 32x32 -> 16x16 -> 8x8 -> 4x4 -> 2x2 -> 1x1 
        # The final feature map size is 1x1 with 512 channels.
        
        # Classifier Block
        self.classifier = nn.Sequential(
            nn.Linear(512 * 1 * 1, 4096),  # Final feature size is 512x1x1
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, 4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096, num_classes)
        )
        
        # Initialize weights (Crucial for VGG to train well)
        self._initialize_weights()

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, 1) 
        logits = self.classifier(x)
        return logits
    
    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')  # Kaiming init standard for ReLU
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

# Helper function to create the feature extractor based on VGG16_arch
def make_layers(cfg, batch_norm=False):
    layers = []
    in_channels = 3 # Start with 3 color channels
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

# Instantiate the Model 
set_all_seeds(RANDOM_SEED)
vgg16_features = make_layers(VGG16_arch, batch_norm=True) # Use Batch Norm for better convergence
model = VGG(vgg16_features, num_classes=NUM_CLASSES).to(device)


# Print VGG16 architecture
print(vgg16_features)


# Print Summary
print("--- VGG-16 Model Summary ---")
summary(model, 
        input_size=(BATCH_SIZE, 3, INPUT_SIZE, INPUT_SIZE), 
        col_names=["input_size", "output_size", "num_params", "kernel_size", "mult_adds"]
       )

Sequential(
  (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (2): ReLU(inplace=True)
  (3): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (4): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (5): ReLU(inplace=True)
  (6): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (7): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (8): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (9): ReLU(inplace=True)
  (10): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (11): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (12): ReLU(inplace=True)
  (13): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (14): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (15): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (16): ReLU(inplace=True)
  (17): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (18): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (19): ReLU(inplace=True)
  (20): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (21): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (22): ReLU(inplace=True)
  (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (24): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (25): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (26): ReLU(inplace=True)
  (27): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (28): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (29): ReLU(inplace=True)
  (30): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (31): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (32): ReLU(inplace=True)
  (33): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (34): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (35): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (36): ReLU(inplace=True)
  (37): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (38): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (39): ReLU(inplace=True)
  (40): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (41): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
  (42): ReLU(inplace=True)
  (43): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
--- VGG-16 Model Summary ---

=====================================================================================================================================================================
Layer (type:depth-idx)                   Input Shape               Output Shape              Param #                   Kernel Shape              Mult-Adds
=====================================================================================================================================================================
VGG                                      [128, 3, 32, 32]          [128, 10]                 --                        --                        --
├─Sequential: 1-1                        [128, 3, 32, 32]          [128, 512, 1, 1]          --                        --                        --
│    └─Conv2d: 2-1                       [128, 3, 32, 32]          [128, 64, 32, 32]         1,792                     [3, 3]                    234,881,024
│    └─BatchNorm2d: 2-2                  [128, 64, 32, 32]         [128, 64, 32, 32]         128                       --                        16,384
│    └─ReLU: 2-3                         [128, 64, 32, 32]         [128, 64, 32, 32]         --                        --                        --
│    └─Conv2d: 2-4                       [128, 64, 32, 32]         [128, 64, 32, 32]         36,928                    [3, 3]                    4,840,226,816
│    └─BatchNorm2d: 2-5                  [128, 64, 32, 32]         [128, 64, 32, 32]         128                       --                        16,384
│    └─ReLU: 2-6                         [128, 64, 32, 32]         [128, 64, 32, 32]         --                        --                        --
│    └─MaxPool2d: 2-7                    [128, 64, 32, 32]         [128, 64, 16, 16]         --                        2                         --
│    └─Conv2d: 2-8                       [128, 64, 16, 16]         [128, 128, 16, 16]        73,856                    [3, 3]                    2,420,113,408
│    └─BatchNorm2d: 2-9                  [128, 128, 16, 16]        [128, 128, 16, 16]        256                       --                        32,768
│    └─ReLU: 2-10                        [128, 128, 16, 16]        [128, 128, 16, 16]        --                        --                        --
│    └─Conv2d: 2-11                      [128, 128, 16, 16]        [128, 128, 16, 16]        147,584                   [3, 3]                    4,836,032,512
│    └─BatchNorm2d: 2-12                 [128, 128, 16, 16]        [128, 128, 16, 16]        256                       --                        32,768
│    └─ReLU: 2-13                        [128, 128, 16, 16]        [128, 128, 16, 16]        --                        --                        --
│    └─MaxPool2d: 2-14                   [128, 128, 16, 16]        [128, 128, 8, 8]          --                        2                         --
│    └─Conv2d: 2-15                      [128, 128, 8, 8]          [128, 256, 8, 8]          295,168                   [3, 3]                    2,418,016,256
│    └─BatchNorm2d: 2-16                 [128, 256, 8, 8]          [128, 256, 8, 8]          512                       --                        65,536
│    └─ReLU: 2-17                        [128, 256, 8, 8]          [128, 256, 8, 8]          --                        --                        --
│    └─Conv2d: 2-18                      [128, 256, 8, 8]          [128, 256, 8, 8]          590,080                   [3, 3]                    4,833,935,360
│    └─BatchNorm2d: 2-19                 [128, 256, 8, 8]          [128, 256, 8, 8]          512                       --                        65,536
│    └─ReLU: 2-20                        [128, 256, 8, 8]          [128, 256, 8, 8]          --                        --                        --
│    └─Conv2d: 2-21                      [128, 256, 8, 8]          [128, 256, 8, 8]          590,080                   [3, 3]                    4,833,935,360
│    └─BatchNorm2d: 2-22                 [128, 256, 8, 8]          [128, 256, 8, 8]          512                       --                        65,536
│    └─ReLU: 2-23                        [128, 256, 8, 8]          [128, 256, 8, 8]          --                        --                        --
│    └─MaxPool2d: 2-24                   [128, 256, 8, 8]          [128, 256, 4, 4]          --                        2                         --
│    └─Conv2d: 2-25                      [128, 256, 4, 4]          [128, 512, 4, 4]          1,180,160                 [3, 3]                    2,416,967,680
│    └─BatchNorm2d: 2-26                 [128, 512, 4, 4]          [128, 512, 4, 4]          1,024                     --                        131,072
│    └─ReLU: 2-27                        [128, 512, 4, 4]          [128, 512, 4, 4]          --                        --                        --
│    └─Conv2d: 2-28                      [128, 512, 4, 4]          [128, 512, 4, 4]          2,359,808                 [3, 3]                    4,832,886,784
│    └─BatchNorm2d: 2-29                 [128, 512, 4, 4]          [128, 512, 4, 4]          1,024                     --                        131,072
│    └─ReLU: 2-30                        [128, 512, 4, 4]          [128, 512, 4, 4]          --                        --                        --
│    └─Conv2d: 2-31                      [128, 512, 4, 4]          [128, 512, 4, 4]          2,359,808                 [3, 3]                    4,832,886,784
│    └─BatchNorm2d: 2-32                 [128, 512, 4, 4]          [128, 512, 4, 4]          1,024                     --                        131,072
│    └─ReLU: 2-33                        [128, 512, 4, 4]          [128, 512, 4, 4]          --                        --                        --
│    └─MaxPool2d: 2-34                   [128, 512, 4, 4]          [128, 512, 2, 2]          --                        2                         --
│    └─Conv2d: 2-35                      [128, 512, 2, 2]          [128, 512, 2, 2]          2,359,808                 [3, 3]                    1,208,221,696
│    └─BatchNorm2d: 2-36                 [128, 512, 2, 2]          [128, 512, 2, 2]          1,024                     --                        131,072
│    └─ReLU: 2-37                        [128, 512, 2, 2]          [128, 512, 2, 2]          --                        --                        --
│    └─Conv2d: 2-38                      [128, 512, 2, 2]          [128, 512, 2, 2]          2,359,808                 [3, 3]                    1,208,221,696
│    └─BatchNorm2d: 2-39                 [128, 512, 2, 2]          [128, 512, 2, 2]          1,024                     --                        131,072
│    └─ReLU: 2-40                        [128, 512, 2, 2]          [128, 512, 2, 2]          --                        --                        --
│    └─Conv2d: 2-41                      [128, 512, 2, 2]          [128, 512, 2, 2]          2,359,808                 [3, 3]                    1,208,221,696
│    └─BatchNorm2d: 2-42                 [128, 512, 2, 2]          [128, 512, 2, 2]          1,024                     --                        131,072
│    └─ReLU: 2-43                        [128, 512, 2, 2]          [128, 512, 2, 2]          --                        --                        --
│    └─MaxPool2d: 2-44                   [128, 512, 2, 2]          [128, 512, 1, 1]          --                        2                         --
├─Sequential: 1-2                        [128, 512]                [128, 10]                 --                        --                        --
│    └─Linear: 2-45                      [128, 512]                [128, 4096]               2,101,248                 --                        268,959,744
│    └─ReLU: 2-46                        [128, 4096]               [128, 4096]               --                        --                        --
│    └─Dropout: 2-47                     [128, 4096]               [128, 4096]               --                        --                        --
│    └─Linear: 2-48                      [128, 4096]               [128, 4096]               16,781,312                --                        2,148,007,936
│    └─ReLU: 2-49                        [128, 4096]               [128, 4096]               --                        --                        --
│    └─Dropout: 2-50                     [128, 4096]               [128, 4096]               --                        --                        --
│    └─Linear: 2-51                      [128, 4096]               [128, 10]                 40,970                    --                        5,244,160
=====================================================================================================================================================================
Total params: 33,646,666
Trainable params: 33,646,666
Non-trainable params: 0
Total mult-adds (Units.GIGABYTES): 42.55
=====================================================================================================================================================================
Input size (MB): 1.57
Forward/backward pass size (MB): 574.63
Params size (MB): 134.59
Estimated Total Size (MB): 710.79
=====================================================================================================================================================================

---

3. Model Training, Testing & Evaluation

---

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=5e-4)
# optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=0.9, weight_decay=5e-4) # Use weight decay (L2) to combat overfitting
# scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5)

def calculate_accuracy(loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return 100 * correct / total

# Training Loop
start_time = time.time()
history = {'train_loss': [], 'val_loss': [], 'train_acc': [], 'val_acc': [], 'lrs': []}

print(f"Starting training on {device}...")

for epoch in range(NUM_EPOCHS):
    model.train()
    running_loss = 0.0
    
    for i, (inputs, labels) in enumerate(train_loader):
        inputs, labels = inputs.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item() * inputs.size(0)

    train_loss = running_loss / len(train_subset)
    train_acc = calculate_accuracy(train_loader)
    
    # Calculate validation loss and accuracy
    model.eval()
    val_loss = 0.0
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            val_loss += loss.item() * inputs.size(0)
    
    val_loss /= len(val_subset)
    val_acc = calculate_accuracy(val_loader)
    
    # scheduler.step(val_loss)
    # current_lr = optimizer.param_groups[0]['lr']

    # Log metrics
    history['train_loss'].append(train_loss)
    history['val_loss'].append(val_loss)
    history['train_acc'].append(train_acc)
    history['val_acc'].append(val_acc)
    # history['lrs'].append(current_lr)

    print(f'Epoch {epoch + 1}/{NUM_EPOCHS} | '
          f'Train Loss: {train_loss:.5f}, '
          f'Train Acc.: {train_acc:.2f}% | '
          f'Val Loss: {val_loss:.5f}, '
          f'Val Acc.: {val_acc:.2f}% | '
          # f'LR: {current_lr:.8f}')
         )

print('\nTotal Training Time: %.2f min' % ((time.time() - start_time)/60))

# --- Final Evaluation ---
test_accuracy = calculate_accuracy(test_loader)
print(f'\nTest Accuracy: {test_accuracy:.2f}%')

Starting training on cuda...
Epoch 1/20 | Train Loss: 1.95216, Train Acc.: 26.73% | Val Loss: 1.86548, Val Acc.: 25.84% | 
Epoch 2/20 | Train Loss: 1.66956, Train Acc.: 34.86% | Val Loss: 1.56075, Val Acc.: 35.00% | 
Epoch 3/20 | Train Loss: 1.50966, Train Acc.: 39.43% | Val Loss: 1.51754, Val Acc.: 39.80% | 
Epoch 4/20 | Train Loss: 1.41668, Train Acc.: 42.27% | Val Loss: 1.44585, Val Acc.: 40.96% | 
Epoch 5/20 | Train Loss: 1.29995, Train Acc.: 51.18% | Val Loss: 1.28786, Val Acc.: 50.78% | 
Epoch 6/20 | Train Loss: 1.11770, Train Acc.: 55.07% | Val Loss: 1.31134, Val Acc.: 54.18% | 
Epoch 7/20 | Train Loss: 0.96871, Train Acc.: 65.46% | Val Loss: 0.99572, Val Acc.: 65.04% | 
Epoch 8/20 | Train Loss: 0.85259, Train Acc.: 68.81% | Val Loss: 0.95486, Val Acc.: 69.14% | 
Epoch 9/20 | Train Loss: 0.78167, Train Acc.: 71.54% | Val Loss: 0.86358, Val Acc.: 70.28% | 
Epoch 10/20 | Train Loss: 0.72179, Train Acc.: 75.38% | Val Loss: 0.80232, Val Acc.: 73.72% | 
Epoch 11/20 | Train Loss: 0.65424, Train Acc.: 78.11% | Val Loss: 0.73177, Val Acc.: 76.94% | 
Epoch 12/20 | Train Loss: 0.61499, Train Acc.: 75.81% | Val Loss: 0.81416, Val Acc.: 74.34% | 
Epoch 13/20 | Train Loss: 0.58261, Train Acc.: 83.34% | Val Loss: 0.57393, Val Acc.: 81.04% | 
Epoch 14/20 | Train Loss: 0.55062, Train Acc.: 79.63% | Val Loss: 0.70647, Val Acc.: 77.66% | 
Epoch 15/20 | Train Loss: 0.52555, Train Acc.: 81.58% | Val Loss: 0.72878, Val Acc.: 78.58% | 
Epoch 16/20 | Train Loss: 0.50517, Train Acc.: 83.65% | Val Loss: 0.59254, Val Acc.: 81.28% | 
Epoch 17/20 | Train Loss: 0.49256, Train Acc.: 83.88% | Val Loss: 0.60581, Val Acc.: 80.08% | 
Epoch 18/20 | Train Loss: 0.46611, Train Acc.: 87.10% | Val Loss: 0.51282, Val Acc.: 83.68% | 
Epoch 19/20 | Train Loss: 0.45528, Train Acc.: 87.18% | Val Loss: 0.52089, Val Acc.: 83.74% | 
Epoch 20/20 | Train Loss: 0.44616, Train Acc.: 87.14% | Val Loss: 0.52678, Val Acc.: 83.18% | 

Total Training Time: 21.36 min

Test Accuracy: 84.05%

def plot_history(history):
    """
    Plots the training and validation loss, accuracy, and learning rate
    from a PyTorch history dictionary.

    Args:
        history (dict): Dictionary containing training metrics.
                        Expected keys: 'train_loss', 'val_loss', 'train_acc', 'val_acc', 'lrs'.
    """
    epochs = range(len(history['train_loss']))

    plt.figure(figsize=(18, 5))

    # Plot 1: Loss
    plt.subplot(1, 3, 1)
    plt.plot(epochs, history['train_loss'], label='Training Loss')
    plt.plot(epochs, history['val_loss'], label='Validation Loss')
    plt.title('Loss')
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    # plt.grid(True)

    # Plot 2: Accuracy
    plt.subplot(1, 3, 2)
    plt.plot(epochs, history['train_acc'], label='Training Accuracy')
    plt.plot(epochs, history['val_acc'], label='Validation Accuracy')
    plt.title('Accuracy')
    plt.xlabel('Epochs')
    plt.ylabel('Accuracy')
    plt.legend()
    # plt.grid(True)

    plt.tight_layout()
    plt.show() 
    # plt.savefig() 

plot_history(history)