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
In [1]:
!pip install watermark
Requirement already satisfied: watermark in /usr/local/lib/python3.12/dist-packages (2.6.0) Requirement already satisfied: ipython>=6.0 in /usr/local/lib/python3.12/dist-packages (from watermark) (7.34.0) Requirement already satisfied: importlib-metadata>=1.4 in /usr/local/lib/python3.12/dist-packages (from watermark) (8.7.0) Requirement already satisfied: setuptools in /usr/local/lib/python3.12/dist-packages (from watermark) (75.2.0) Requirement already satisfied: zipp>=3.20 in /usr/local/lib/python3.12/dist-packages (from importlib-metadata>=1.4->watermark) (3.23.0) Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (0.19.2) Requirement already satisfied: decorator in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (4.4.2) Requirement already satisfied: pickleshare in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (0.7.5) Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.12/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.12/dist-packages (from ipython>=6.0->watermark) (3.0.52) Requirement already satisfied: pygments in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (2.19.2) Requirement already satisfied: backcall in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (0.2.0) Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.12/dist-packages (from ipython>=6.0->watermark) (0.2.1) Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.12/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.12/dist-packages (from jedi>=0.16->ipython>=6.0->watermark) (0.8.5) Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.12/dist-packages (from pexpect>4.3->ipython>=6.0->watermark) (0.7.0) Requirement already satisfied: wcwidth in /usr/local/lib/python3.12/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython>=6.0->watermark) (0.2.14)
In [2]:
!pip install torchinfo
Requirement already satisfied: torchinfo in /usr/local/lib/python3.12/dist-packages (1.8.0)
In [3]:
%load_ext watermark
# %reload_ext watermark
%watermark -a 'Chizoba Obasi' -v -p torch
Author: Chizoba Obasi Python implementation: CPython Python version : 3.12.12 IPython version : 7.34.0 torch: 2.9.0+cu126
MobileNetV3-Large Classifier on CIFAR-10¶
This notebook implements MobileNetV3-Large [1] and applies it to the CIFAR-10 dataset.
Network Architecture¶
MobileNet-V3: Bottleneck Block¶
MobileNet-Vx: Backbone Block¶
| MobileNetV2 | MobileNetV3 |
|---|---|
 Source |
 Source |
REFERENCES:
- Howard, M., et al. "Searching for MobileNetV3." CVPR, 2018.
- Johnson, J. (2022, February 2). Lecture 11: CNN Architectures, Part 2 [PowerPoint Slides]. University of Michigan EECS 498/598: Deep Learning for Computer Vision. pp 44-53. https://web.eecs.umich.edu/~justincj/slides/eecs498/WI2022/598_WI2022_lecture11.pdf
In [4]:
import torch
import torchvision
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
from torchinfo import summary
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
In [5]:
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)
In [6]:
# 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: 1 GPU 0: Tesla T4
In [7]:
##########################
### SETTINGS
##########################
# Hyperparameters
RANDOM_SEED = 123
BATCH_SIZE = 128 #256
LEARNING_RATE = 0.001
NUM_EPOCHS = 75 #50 #30
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
In [8]:
set_all_seeds(RANDOM_SEED)
1. Data Preparation & Loading¶
Notes:
- Since we adjust the model to handle the CIFAR images natively, we should remove the
Resize((xx, yy))from your transforms. It will make ythe training 10x faster and use significantly less memory.
In [9]:
# CIFAR-10 specific stats (MEAN_arr, STD_arr)
stats = ((0.4914, 0.4822, 0.4465), (0.2471, 0.2435, 0.2616)) #(0.2023, 0.1994, 0.2010))
train_transforms = transforms.Compose([
transforms.RandomCrop(32, padding=4), # Standard CIFAR augmentation
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(*stats)
])
test_transforms = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(*stats)
])
# 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, num_workers=2, pin_memory=True)
val_loader = DataLoader(val_subset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
In [10]:
# Checking the dataset
for images, labels in train_loader:
print('Image batch dimensions:', images.shape)
print('Image label dimensions:', labels.shape)
print('Class labels of 10 examples:', labels[:10])
break
Image batch dimensions: torch.Size([128, 3, 32, 32]) Image label dimensions: torch.Size([128]) Class labels of 10 examples: tensor([2, 7, 4, 7, 0, 1, 1, 6, 9, 2])
In [11]:
plt.figure(figsize=(8, 8))
plt.axis("off")
plt.title("Training Images")
plt.imshow(np.transpose(torchvision.utils.make_grid(images[:64],
padding=2, normalize=True),
(1, 2, 0)))
Out[11]:
<matplotlib.image.AxesImage at 0x7928e4015370>
2. Model Creation¶
Notes:
- For CIFAR-10, we modify the original architecture by setting the initial convolution stride to 1 instead of 2. This prevents the $32 \times 32$ image from shrinking too fast.
In [12]:
def _make_divisible(v, divisor, min_value=None):
""" Ensure all layers have a channel count divisible by 8 """
if min_value is None:
min_value = divisor
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
if new_v < 0.9 * v:
new_v += divisor
return new_v
class HSwish(nn.Module):
"""
Hard-Swish activation: professional fallback implementation.
"""
def __init__(self, inplace: bool = True):
super(HSwish, self).__init__()
self.inplace = inplace
def forward(self, x: torch.Tensor) -> torch.Tensor:
if hasattr(F, 'hardswish'):
return F.hardswish(x, inplace=self.inplace)
return x * F.relu6(x + 3, inplace=self.inplace) / 6
class SqueezeExcitation(nn.Module):
"""
SE Block: Learns to weight channels based on their importance.
"""
def __init__(self, in_channels, reduction=4):
super(SqueezeExcitation, self).__init__()
# Ensure channel counts are divisible by 8 for hardware efficiency
squeeze_channels = _make_divisible(in_channels // reduction, 8)
self.se = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Conv2d(in_channels, squeeze_channels, kernel_size=1),
nn.ReLU(inplace=True),
nn.Conv2d(squeeze_channels, in_channels, kernel_size=1),
nn.Hardsigmoid(inplace=True)
)
def forward(self, x):
return x * self.se(x)
class Block(nn.Module):
def __init__(self, kernel_size, in_channels, expand_size, out_channels, se, nl, stride):
super().__init__()
self.stride = stride
self.use_se = se
self.use_res_connect = self.stride == 1 and in_channels == out_channels
activation = HSwish if nl == 'HS' else nn.ReLU
# Expansion
self.conv1 = nn.Conv2d(in_channels, expand_size, 1, 1, 0, bias=False)
self.bn1 = nn.BatchNorm2d(expand_size)
self.act1 = activation(inplace=True)
# Depthwise
self.conv2 = nn.Conv2d(expand_size, expand_size, kernel_size, stride, kernel_size//2, groups=expand_size, bias=False)
self.bn2 = nn.BatchNorm2d(expand_size)
self.act2 = activation(inplace=True)
# SE Block
self.se = SqueezeExcitation(expand_size) if se else nn.Identity()
# Projection (Linear)
self.conv3 = nn.Conv2d(expand_size, out_channels, 1, 1, 0, bias=False)
self.bn3 = nn.BatchNorm2d(out_channels)
def forward(self, x):
out = self.act1(self.bn1(self.conv1(x)))
out = self.act2(self.bn2(self.conv2(out)))
out = self.se(out)
out = self.bn3(self.conv3(out))
if self.use_res_connect:
return x + out
return out
class MobileNetV3_Large(nn.Module):
def __init__(self, num_classes=10):
super().__init__()
# Configuration for Large: [kernel, expansion, out_channels, use_se, activation, stride]
# Based on the official MobileNetV3 paper
self.config = [
[3, 16, 16, False, 'RE', 1],
[3, 64, 24, False, 'RE', 2],
[3, 72, 24, False, 'RE', 1],
[5, 72, 40, True, 'RE', 2],
[5, 120, 40, True, 'RE', 1],
[5, 120, 40, True, 'RE', 1],
[3, 240, 80, False, 'HS', 2],
[3, 200, 80, False, 'HS', 1],
[3, 184, 80, False, 'HS', 1],
[3, 184, 80, False, 'HS', 1],
[3, 480, 112, True, 'HS', 1],
[3, 672, 112, True, 'HS', 1],
[5, 672, 160, True, 'HS', 2],
[5, 960, 160, True, 'HS', 1],
[5, 960, 160, True, 'HS', 1],
]
# First Layer: Stride 1 for CIFAR-10 (originally 2 in ImageNet)
self.conv1 = nn.Conv2d(3, 16, 3, 1, 1, bias=False)
self.bn1 = nn.BatchNorm2d(16)
self.act1 = HSwish(inplace=True)
# Build Blocks
layers = []
in_channels = 16
for k, exp, out, se, nl, s in self.config:
# Use _make_divisible for hardware-optimized channel counts
exp_ch = _make_divisible(exp, 8)
out_ch = _make_divisible(out, 8)
layers.append(Block(k, in_channels, exp_ch, out_ch, se, nl, s))
in_channels = out_ch
self.blocks = nn.Sequential(*layers)
# Final Feature Layers
last_out = 960
self.conv2 = nn.Conv2d(in_channels, last_out, 1, 1, 0, bias=False)
self.bn2 = nn.BatchNorm2d(last_out)
self.act2 = HSwish(inplace=True)
# Global Average Pooling and Classifier
self.avgpool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Sequential(
nn.Linear(last_out, 1280),
HSwish(inplace=True),
nn.Dropout(0.2),
nn.Linear(1280, num_classes)
)
def forward(self, x):
x = self.act1(self.bn1(self.conv1(x)))
x = self.blocks(x)
x = self.act2(self.bn2(self.conv2(x)))
x = self.avgpool(x)
x = x.view(x.size(0), -1)
return self.fc(x)
In [13]:
model = MobileNetV3_Large().to(device)
In [14]:
# Model parameter details
summary(model,
input_size=(BATCH_SIZE, 3, INPUT_SIZE, INPUT_SIZE),
# verbose=1,
col_names=["input_size", "output_size", "num_params", "kernel_size", "mult_adds"]
)
Out[14]:
=============================================================================================================================================================================== Layer (type:depth-idx) Input Shape Output Shape Param # Kernel Shape Mult-Adds =============================================================================================================================================================================== MobileNetV3_Large [128, 3, 32, 32] [128, 10] -- -- -- ├─Conv2d: 1-1 [128, 3, 32, 32] [128, 16, 32, 32] 432 [3, 3] 56,623,104 ├─BatchNorm2d: 1-2 [128, 16, 32, 32] [128, 16, 32, 32] 32 -- 4,096 ├─HSwish: 1-3 [128, 16, 32, 32] [128, 16, 32, 32] -- -- -- ├─Sequential: 1-4 [128, 16, 32, 32] [128, 160, 2, 2] -- -- -- │ └─Block: 2-1 [128, 16, 32, 32] [128, 16, 32, 32] -- -- -- │ │ └─Conv2d: 3-1 [128, 16, 32, 32] [128, 16, 32, 32] 256 [1, 1] 33,554,432 │ │ └─BatchNorm2d: 3-2 [128, 16, 32, 32] [128, 16, 32, 32] 32 -- 4,096 │ │ └─ReLU: 3-3 [128, 16, 32, 32] [128, 16, 32, 32] -- -- -- │ │ └─Conv2d: 3-4 [128, 16, 32, 32] [128, 16, 32, 32] 144 [3, 3] 18,874,368 │ │ └─BatchNorm2d: 3-5 [128, 16, 32, 32] [128, 16, 32, 32] 32 -- 4,096 │ │ └─ReLU: 3-6 [128, 16, 32, 32] [128, 16, 32, 32] -- -- -- │ │ └─Identity: 3-7 [128, 16, 32, 32] [128, 16, 32, 32] -- -- -- │ │ └─Conv2d: 3-8 [128, 16, 32, 32] [128, 16, 32, 32] 256 [1, 1] 33,554,432 │ │ └─BatchNorm2d: 3-9 [128, 16, 32, 32] [128, 16, 32, 32] 32 -- 4,096 │ └─Block: 2-2 [128, 16, 32, 32] [128, 24, 16, 16] -- -- -- │ │ └─Conv2d: 3-10 [128, 16, 32, 32] [128, 64, 32, 32] 1,024 [1, 1] 134,217,728 │ │ └─BatchNorm2d: 3-11 [128, 64, 32, 32] [128, 64, 32, 32] 128 -- 16,384 │ │ └─ReLU: 3-12 [128, 64, 32, 32] [128, 64, 32, 32] -- -- -- │ │ └─Conv2d: 3-13 [128, 64, 32, 32] [128, 64, 16, 16] 576 [3, 3] 18,874,368 │ │ └─BatchNorm2d: 3-14 [128, 64, 16, 16] [128, 64, 16, 16] 128 -- 16,384 │ │ └─ReLU: 3-15 [128, 64, 16, 16] [128, 64, 16, 16] -- -- -- │ │ └─Identity: 3-16 [128, 64, 16, 16] [128, 64, 16, 16] -- -- -- │ │ └─Conv2d: 3-17 [128, 64, 16, 16] [128, 24, 16, 16] 1,536 [1, 1] 50,331,648 │ │ └─BatchNorm2d: 3-18 [128, 24, 16, 16] [128, 24, 16, 16] 48 -- 6,144 │ └─Block: 2-3 [128, 24, 16, 16] [128, 24, 16, 16] -- -- -- │ │ └─Conv2d: 3-19 [128, 24, 16, 16] [128, 72, 16, 16] 1,728 [1, 1] 56,623,104 │ │ └─BatchNorm2d: 3-20 [128, 72, 16, 16] [128, 72, 16, 16] 144 -- 18,432 │ │ └─ReLU: 3-21 [128, 72, 16, 16] [128, 72, 16, 16] -- -- -- │ │ └─Conv2d: 3-22 [128, 72, 16, 16] [128, 72, 16, 16] 648 [3, 3] 21,233,664 │ │ └─BatchNorm2d: 3-23 [128, 72, 16, 16] [128, 72, 16, 16] 144 -- 18,432 │ │ └─ReLU: 3-24 [128, 72, 16, 16] [128, 72, 16, 16] -- -- -- │ │ └─Identity: 3-25 [128, 72, 16, 16] [128, 72, 16, 16] -- -- -- │ │ └─Conv2d: 3-26 [128, 72, 16, 16] [128, 24, 16, 16] 1,728 [1, 1] 56,623,104 │ │ └─BatchNorm2d: 3-27 [128, 24, 16, 16] [128, 24, 16, 16] 48 -- 6,144 │ └─Block: 2-4 [128, 24, 16, 16] [128, 40, 8, 8] -- -- -- │ │ └─Conv2d: 3-28 [128, 24, 16, 16] [128, 72, 16, 16] 1,728 [1, 1] 56,623,104 │ │ └─BatchNorm2d: 3-29 [128, 72, 16, 16] [128, 72, 16, 16] 144 -- 18,432 │ │ └─ReLU: 3-30 [128, 72, 16, 16] [128, 72, 16, 16] -- -- -- │ │ └─Conv2d: 3-31 [128, 72, 16, 16] [128, 72, 8, 8] 1,800 [5, 5] 14,745,600 │ │ └─BatchNorm2d: 3-32 [128, 72, 8, 8] [128, 72, 8, 8] 144 -- 18,432 │ │ └─ReLU: 3-33 [128, 72, 8, 8] [128, 72, 8, 8] -- -- -- │ │ └─SqueezeExcitation: 3-34 [128, 72, 8, 8] [128, 72, 8, 8] 3,552 -- 454,656 │ │ └─Conv2d: 3-35 [128, 72, 8, 8] [128, 40, 8, 8] 2,880 [1, 1] 23,592,960 │ │ └─BatchNorm2d: 3-36 [128, 40, 8, 8] [128, 40, 8, 8] 80 -- 10,240 │ └─Block: 2-5 [128, 40, 8, 8] [128, 40, 8, 8] -- -- -- │ │ └─Conv2d: 3-37 [128, 40, 8, 8] [128, 120, 8, 8] 4,800 [1, 1] 39,321,600 │ │ └─BatchNorm2d: 3-38 [128, 120, 8, 8] [128, 120, 8, 8] 240 -- 30,720 │ │ └─ReLU: 3-39 [128, 120, 8, 8] [128, 120, 8, 8] -- -- -- │ │ └─Conv2d: 3-40 [128, 120, 8, 8] [128, 120, 8, 8] 3,000 [5, 5] 24,576,000 │ │ └─BatchNorm2d: 3-41 [128, 120, 8, 8] [128, 120, 8, 8] 240 -- 30,720 │ │ └─ReLU: 3-42 [128, 120, 8, 8] [128, 120, 8, 8] -- -- -- │ │ └─SqueezeExcitation: 3-43 [128, 120, 8, 8] [128, 120, 8, 8] 7,832 -- 1,002,496 │ │ └─Conv2d: 3-44 [128, 120, 8, 8] [128, 40, 8, 8] 4,800 [1, 1] 39,321,600 │ │ └─BatchNorm2d: 3-45 [128, 40, 8, 8] [128, 40, 8, 8] 80 -- 10,240 │ └─Block: 2-6 [128, 40, 8, 8] [128, 40, 8, 8] -- -- -- │ │ └─Conv2d: 3-46 [128, 40, 8, 8] [128, 120, 8, 8] 4,800 [1, 1] 39,321,600 │ │ └─BatchNorm2d: 3-47 [128, 120, 8, 8] [128, 120, 8, 8] 240 -- 30,720 │ │ └─ReLU: 3-48 [128, 120, 8, 8] [128, 120, 8, 8] -- -- -- │ │ └─Conv2d: 3-49 [128, 120, 8, 8] [128, 120, 8, 8] 3,000 [5, 5] 24,576,000 │ │ └─BatchNorm2d: 3-50 [128, 120, 8, 8] [128, 120, 8, 8] 240 -- 30,720 │ │ └─ReLU: 3-51 [128, 120, 8, 8] [128, 120, 8, 8] -- -- -- │ │ └─SqueezeExcitation: 3-52 [128, 120, 8, 8] [128, 120, 8, 8] 7,832 -- 1,002,496 │ │ └─Conv2d: 3-53 [128, 120, 8, 8] [128, 40, 8, 8] 4,800 [1, 1] 39,321,600 │ │ └─BatchNorm2d: 3-54 [128, 40, 8, 8] [128, 40, 8, 8] 80 -- 10,240 │ └─Block: 2-7 [128, 40, 8, 8] [128, 80, 4, 4] -- -- -- │ │ └─Conv2d: 3-55 [128, 40, 8, 8] [128, 240, 8, 8] 9,600 [1, 1] 78,643,200 │ │ └─BatchNorm2d: 3-56 [128, 240, 8, 8] [128, 240, 8, 8] 480 -- 61,440 │ │ └─HSwish: 3-57 [128, 240, 8, 8] [128, 240, 8, 8] -- -- -- │ │ └─Conv2d: 3-58 [128, 240, 8, 8] [128, 240, 4, 4] 2,160 [3, 3] 4,423,680 │ │ └─BatchNorm2d: 3-59 [128, 240, 4, 4] [128, 240, 4, 4] 480 -- 61,440 │ │ └─HSwish: 3-60 [128, 240, 4, 4] [128, 240, 4, 4] -- -- -- │ │ └─Identity: 3-61 [128, 240, 4, 4] [128, 240, 4, 4] -- -- -- │ │ └─Conv2d: 3-62 [128, 240, 4, 4] [128, 80, 4, 4] 19,200 [1, 1] 39,321,600 │ │ └─BatchNorm2d: 3-63 [128, 80, 4, 4] [128, 80, 4, 4] 160 -- 20,480 │ └─Block: 2-8 [128, 80, 4, 4] [128, 80, 4, 4] -- -- -- │ │ └─Conv2d: 3-64 [128, 80, 4, 4] [128, 200, 4, 4] 16,000 [1, 1] 32,768,000 │ │ └─BatchNorm2d: 3-65 [128, 200, 4, 4] [128, 200, 4, 4] 400 -- 51,200 │ │ └─HSwish: 3-66 [128, 200, 4, 4] [128, 200, 4, 4] -- -- -- │ │ └─Conv2d: 3-67 [128, 200, 4, 4] [128, 200, 4, 4] 1,800 [3, 3] 3,686,400 │ │ └─BatchNorm2d: 3-68 [128, 200, 4, 4] [128, 200, 4, 4] 400 -- 51,200 │ │ └─HSwish: 3-69 [128, 200, 4, 4] [128, 200, 4, 4] -- -- -- │ │ └─Identity: 3-70 [128, 200, 4, 4] [128, 200, 4, 4] -- -- -- │ │ └─Conv2d: 3-71 [128, 200, 4, 4] [128, 80, 4, 4] 16,000 [1, 1] 32,768,000 │ │ └─BatchNorm2d: 3-72 [128, 80, 4, 4] [128, 80, 4, 4] 160 -- 20,480 │ └─Block: 2-9 [128, 80, 4, 4] [128, 80, 4, 4] -- -- -- │ │ └─Conv2d: 3-73 [128, 80, 4, 4] [128, 184, 4, 4] 14,720 [1, 1] 30,146,560 │ │ └─BatchNorm2d: 3-74 [128, 184, 4, 4] [128, 184, 4, 4] 368 -- 47,104 │ │ └─HSwish: 3-75 [128, 184, 4, 4] [128, 184, 4, 4] -- -- -- │ │ └─Conv2d: 3-76 [128, 184, 4, 4] [128, 184, 4, 4] 1,656 [3, 3] 3,391,488 │ │ └─BatchNorm2d: 3-77 [128, 184, 4, 4] [128, 184, 4, 4] 368 -- 47,104 │ │ └─HSwish: 3-78 [128, 184, 4, 4] [128, 184, 4, 4] -- -- -- │ │ └─Identity: 3-79 [128, 184, 4, 4] [128, 184, 4, 4] -- -- -- │ │ └─Conv2d: 3-80 [128, 184, 4, 4] [128, 80, 4, 4] 14,720 [1, 1] 30,146,560 │ │ └─BatchNorm2d: 3-81 [128, 80, 4, 4] [128, 80, 4, 4] 160 -- 20,480 │ └─Block: 2-10 [128, 80, 4, 4] [128, 80, 4, 4] -- -- -- │ │ └─Conv2d: 3-82 [128, 80, 4, 4] [128, 184, 4, 4] 14,720 [1, 1] 30,146,560 │ │ └─BatchNorm2d: 3-83 [128, 184, 4, 4] [128, 184, 4, 4] 368 -- 47,104 │ │ └─HSwish: 3-84 [128, 184, 4, 4] [128, 184, 4, 4] -- -- -- │ │ └─Conv2d: 3-85 [128, 184, 4, 4] [128, 184, 4, 4] 1,656 [3, 3] 3,391,488 │ │ └─BatchNorm2d: 3-86 [128, 184, 4, 4] [128, 184, 4, 4] 368 -- 47,104 │ │ └─HSwish: 3-87 [128, 184, 4, 4] [128, 184, 4, 4] -- -- -- │ │ └─Identity: 3-88 [128, 184, 4, 4] [128, 184, 4, 4] -- -- -- │ │ └─Conv2d: 3-89 [128, 184, 4, 4] [128, 80, 4, 4] 14,720 [1, 1] 30,146,560 │ │ └─BatchNorm2d: 3-90 [128, 80, 4, 4] [128, 80, 4, 4] 160 -- 20,480 │ └─Block: 2-11 [128, 80, 4, 4] [128, 112, 4, 4] -- -- -- │ │ └─Conv2d: 3-91 [128, 80, 4, 4] [128, 480, 4, 4] 38,400 [1, 1] 78,643,200 │ │ └─BatchNorm2d: 3-92 [128, 480, 4, 4] [128, 480, 4, 4] 960 -- 122,880 │ │ └─HSwish: 3-93 [128, 480, 4, 4] [128, 480, 4, 4] -- -- -- │ │ └─Conv2d: 3-94 [128, 480, 4, 4] [128, 480, 4, 4] 4,320 [3, 3] 8,847,360 │ │ └─BatchNorm2d: 3-95 [128, 480, 4, 4] [128, 480, 4, 4] 960 -- 122,880 │ │ └─HSwish: 3-96 [128, 480, 4, 4] [128, 480, 4, 4] -- -- -- │ │ └─SqueezeExcitation: 3-97 [128, 480, 4, 4] [128, 480, 4, 4] 115,800 -- 14,822,400 │ │ └─Conv2d: 3-98 [128, 480, 4, 4] [128, 112, 4, 4] 53,760 [1, 1] 110,100,480 │ │ └─BatchNorm2d: 3-99 [128, 112, 4, 4] [128, 112, 4, 4] 224 -- 28,672 │ └─Block: 2-12 [128, 112, 4, 4] [128, 112, 4, 4] -- -- -- │ │ └─Conv2d: 3-100 [128, 112, 4, 4] [128, 672, 4, 4] 75,264 [1, 1] 154,140,672 │ │ └─BatchNorm2d: 3-101 [128, 672, 4, 4] [128, 672, 4, 4] 1,344 -- 172,032 │ │ └─HSwish: 3-102 [128, 672, 4, 4] [128, 672, 4, 4] -- -- -- │ │ └─Conv2d: 3-103 [128, 672, 4, 4] [128, 672, 4, 4] 6,048 [3, 3] 12,386,304 │ │ └─BatchNorm2d: 3-104 [128, 672, 4, 4] [128, 672, 4, 4] 1,344 -- 172,032 │ │ └─HSwish: 3-105 [128, 672, 4, 4] [128, 672, 4, 4] -- -- -- │ │ └─SqueezeExcitation: 3-106 [128, 672, 4, 4] [128, 672, 4, 4] 226,632 -- 29,008,896 │ │ └─Conv2d: 3-107 [128, 672, 4, 4] [128, 112, 4, 4] 75,264 [1, 1] 154,140,672 │ │ └─BatchNorm2d: 3-108 [128, 112, 4, 4] [128, 112, 4, 4] 224 -- 28,672 │ └─Block: 2-13 [128, 112, 4, 4] [128, 160, 2, 2] -- -- -- │ │ └─Conv2d: 3-109 [128, 112, 4, 4] [128, 672, 4, 4] 75,264 [1, 1] 154,140,672 │ │ └─BatchNorm2d: 3-110 [128, 672, 4, 4] [128, 672, 4, 4] 1,344 -- 172,032 │ │ └─HSwish: 3-111 [128, 672, 4, 4] [128, 672, 4, 4] -- -- -- │ │ └─Conv2d: 3-112 [128, 672, 4, 4] [128, 672, 2, 2] 16,800 [5, 5] 8,601,600 │ │ └─BatchNorm2d: 3-113 [128, 672, 2, 2] [128, 672, 2, 2] 1,344 -- 172,032 │ │ └─HSwish: 3-114 [128, 672, 2, 2] [128, 672, 2, 2] -- -- -- │ │ └─SqueezeExcitation: 3-115 [128, 672, 2, 2] [128, 672, 2, 2] 226,632 -- 29,008,896 │ │ └─Conv2d: 3-116 [128, 672, 2, 2] [128, 160, 2, 2] 107,520 [1, 1] 55,050,240 │ │ └─BatchNorm2d: 3-117 [128, 160, 2, 2] [128, 160, 2, 2] 320 -- 40,960 │ └─Block: 2-14 [128, 160, 2, 2] [128, 160, 2, 2] -- -- -- │ │ └─Conv2d: 3-118 [128, 160, 2, 2] [128, 960, 2, 2] 153,600 [1, 1] 78,643,200 │ │ └─BatchNorm2d: 3-119 [128, 960, 2, 2] [128, 960, 2, 2] 1,920 -- 245,760 │ │ └─HSwish: 3-120 [128, 960, 2, 2] [128, 960, 2, 2] -- -- -- │ │ └─Conv2d: 3-121 [128, 960, 2, 2] [128, 960, 2, 2] 24,000 [5, 5] 12,288,000 │ │ └─BatchNorm2d: 3-122 [128, 960, 2, 2] [128, 960, 2, 2] 1,920 -- 245,760 │ │ └─HSwish: 3-123 [128, 960, 2, 2] [128, 960, 2, 2] -- -- -- │ │ └─SqueezeExcitation: 3-124 [128, 960, 2, 2] [128, 960, 2, 2] 462,000 -- 59,136,000 │ │ └─Conv2d: 3-125 [128, 960, 2, 2] [128, 160, 2, 2] 153,600 [1, 1] 78,643,200 │ │ └─BatchNorm2d: 3-126 [128, 160, 2, 2] [128, 160, 2, 2] 320 -- 40,960 │ └─Block: 2-15 [128, 160, 2, 2] [128, 160, 2, 2] -- -- -- │ │ └─Conv2d: 3-127 [128, 160, 2, 2] [128, 960, 2, 2] 153,600 [1, 1] 78,643,200 │ │ └─BatchNorm2d: 3-128 [128, 960, 2, 2] [128, 960, 2, 2] 1,920 -- 245,760 │ │ └─HSwish: 3-129 [128, 960, 2, 2] [128, 960, 2, 2] -- -- -- │ │ └─Conv2d: 3-130 [128, 960, 2, 2] [128, 960, 2, 2] 24,000 [5, 5] 12,288,000 │ │ └─BatchNorm2d: 3-131 [128, 960, 2, 2] [128, 960, 2, 2] 1,920 -- 245,760 │ │ └─HSwish: 3-132 [128, 960, 2, 2] [128, 960, 2, 2] -- -- -- │ │ └─SqueezeExcitation: 3-133 [128, 960, 2, 2] [128, 960, 2, 2] 462,000 -- 59,136,000 │ │ └─Conv2d: 3-134 [128, 960, 2, 2] [128, 160, 2, 2] 153,600 [1, 1] 78,643,200 │ │ └─BatchNorm2d: 3-135 [128, 160, 2, 2] [128, 160, 2, 2] 320 -- 40,960 ├─Conv2d: 1-5 [128, 160, 2, 2] [128, 960, 2, 2] 153,600 [1, 1] 78,643,200 ├─BatchNorm2d: 1-6 [128, 960, 2, 2] [128, 960, 2, 2] 1,920 -- 245,760 ├─HSwish: 1-7 [128, 960, 2, 2] [128, 960, 2, 2] -- -- -- ├─AdaptiveAvgPool2d: 1-8 [128, 960, 2, 2] [128, 960, 1, 1] -- -- -- ├─Sequential: 1-9 [128, 960] [128, 10] -- -- -- │ └─Linear: 2-16 [128, 960] [128, 1280] 1,230,080 -- 157,450,240 │ └─HSwish: 2-17 [128, 1280] [128, 1280] -- -- -- │ └─Dropout: 2-18 [128, 1280] [128, 1280] -- -- -- │ └─Linear: 2-19 [128, 1280] [128, 10] 12,810 -- 1,639,680 =============================================================================================================================================================================== Total params: 4,215,130 Trainable params: 4,215,130 Non-trainable params: 0 Total mult-adds (Units.GIGABYTES): 2.61 =============================================================================================================================================================================== Input size (MB): 1.57 Forward/backward pass size (MB): 775.65 Params size (MB): 16.86 Estimated Total Size (MB): 794.09 ===============================================================================================================================================================================
In [15]:
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/75 | Train Loss: 1.74394, Train Acc.: 43.97% | Val Loss: 1.53586, Val Acc.: 42.32% | Epoch 2/75 | Train Loss: 1.38070, Train Acc.: 54.22% | Val Loss: 1.28692, Val Acc.: 53.92% | Epoch 3/75 | Train Loss: 1.18384, Train Acc.: 60.54% | Val Loss: 1.13871, Val Acc.: 58.88% | Epoch 4/75 | Train Loss: 1.04309, Train Acc.: 64.41% | Val Loss: 1.00583, Val Acc.: 62.24% | Epoch 5/75 | Train Loss: 0.94883, Train Acc.: 66.10% | Val Loss: 0.98264, Val Acc.: 64.38% | Epoch 6/75 | Train Loss: 0.86432, Train Acc.: 70.38% | Val Loss: 0.88011, Val Acc.: 69.02% | Epoch 7/75 | Train Loss: 0.80810, Train Acc.: 74.01% | Val Loss: 0.79819, Val Acc.: 72.28% | Epoch 8/75 | Train Loss: 0.75617, Train Acc.: 74.98% | Val Loss: 0.75979, Val Acc.: 73.46% | Epoch 9/75 | Train Loss: 0.71642, Train Acc.: 77.57% | Val Loss: 0.69830, Val Acc.: 75.90% | Epoch 10/75 | Train Loss: 0.68651, Train Acc.: 77.34% | Val Loss: 0.70304, Val Acc.: 75.24% | Epoch 11/75 | Train Loss: 0.65424, Train Acc.: 79.18% | Val Loss: 0.68198, Val Acc.: 76.82% | Epoch 12/75 | Train Loss: 0.62571, Train Acc.: 79.40% | Val Loss: 0.67511, Val Acc.: 77.08% | Epoch 13/75 | Train Loss: 0.60612, Train Acc.: 79.84% | Val Loss: 0.66235, Val Acc.: 76.80% | Epoch 14/75 | Train Loss: 0.59633, Train Acc.: 80.38% | Val Loss: 0.63242, Val Acc.: 77.64% | Epoch 15/75 | Train Loss: 0.57935, Train Acc.: 81.88% | Val Loss: 0.59502, Val Acc.: 79.02% | Epoch 16/75 | Train Loss: 0.55627, Train Acc.: 80.04% | Val Loss: 0.66247, Val Acc.: 76.80% | Epoch 17/75 | Train Loss: 0.54880, Train Acc.: 82.16% | Val Loss: 0.62291, Val Acc.: 79.00% | Epoch 18/75 | Train Loss: 0.53058, Train Acc.: 83.02% | Val Loss: 0.59861, Val Acc.: 79.24% | Epoch 19/75 | Train Loss: 0.52027, Train Acc.: 84.23% | Val Loss: 0.55309, Val Acc.: 80.98% | Epoch 20/75 | Train Loss: 0.50550, Train Acc.: 83.54% | Val Loss: 0.54834, Val Acc.: 80.94% | Epoch 21/75 | Train Loss: 0.50270, Train Acc.: 84.04% | Val Loss: 0.55967, Val Acc.: 80.18% | Epoch 22/75 | Train Loss: 0.49027, Train Acc.: 83.62% | Val Loss: 0.58009, Val Acc.: 80.38% | Epoch 23/75 | Train Loss: 0.48274, Train Acc.: 84.74% | Val Loss: 0.55220, Val Acc.: 81.80% | Epoch 24/75 | Train Loss: 0.47081, Train Acc.: 83.36% | Val Loss: 0.59668, Val Acc.: 79.72% | Epoch 25/75 | Train Loss: 0.46921, Train Acc.: 85.10% | Val Loss: 0.54227, Val Acc.: 81.16% | Epoch 26/75 | Train Loss: 0.45600, Train Acc.: 85.29% | Val Loss: 0.54127, Val Acc.: 82.40% | Epoch 27/75 | Train Loss: 0.44857, Train Acc.: 86.20% | Val Loss: 0.51644, Val Acc.: 82.46% | Epoch 28/75 | Train Loss: 0.44154, Train Acc.: 85.94% | Val Loss: 0.52547, Val Acc.: 82.00% | Epoch 29/75 | Train Loss: 0.44321, Train Acc.: 86.33% | Val Loss: 0.50725, Val Acc.: 82.64% | Epoch 30/75 | Train Loss: 0.43178, Train Acc.: 86.28% | Val Loss: 0.53586, Val Acc.: 81.98% | Epoch 31/75 | Train Loss: 0.42909, Train Acc.: 86.49% | Val Loss: 0.50370, Val Acc.: 83.06% | Epoch 32/75 | Train Loss: 0.42746, Train Acc.: 86.70% | Val Loss: 0.49496, Val Acc.: 82.86% | Epoch 33/75 | Train Loss: 0.41885, Train Acc.: 86.74% | Val Loss: 0.50573, Val Acc.: 83.08% | Epoch 34/75 | Train Loss: 0.41913, Train Acc.: 87.40% | Val Loss: 0.48578, Val Acc.: 82.82% | Epoch 35/75 | Train Loss: 0.41013, Train Acc.: 85.47% | Val Loss: 0.54469, Val Acc.: 81.18% | Epoch 36/75 | Train Loss: 0.40468, Train Acc.: 86.27% | Val Loss: 0.52511, Val Acc.: 82.46% | Epoch 37/75 | Train Loss: 0.40010, Train Acc.: 86.86% | Val Loss: 0.49441, Val Acc.: 82.68% | Epoch 38/75 | Train Loss: 0.39843, Train Acc.: 85.92% | Val Loss: 0.53768, Val Acc.: 82.00% | Epoch 39/75 | Train Loss: 0.39006, Train Acc.: 87.66% | Val Loss: 0.48869, Val Acc.: 83.48% | Epoch 40/75 | Train Loss: 0.38548, Train Acc.: 87.57% | Val Loss: 0.47904, Val Acc.: 83.40% | Epoch 41/75 | Train Loss: 0.38988, Train Acc.: 86.71% | Val Loss: 0.51700, Val Acc.: 81.86% | Epoch 42/75 | Train Loss: 0.38532, Train Acc.: 88.23% | Val Loss: 0.46890, Val Acc.: 84.16% | Epoch 43/75 | Train Loss: 0.37675, Train Acc.: 86.82% | Val Loss: 0.51546, Val Acc.: 82.22% | Epoch 44/75 | Train Loss: 0.37747, Train Acc.: 88.05% | Val Loss: 0.47276, Val Acc.: 84.34% | Epoch 45/75 | Train Loss: 0.37336, Train Acc.: 88.82% | Val Loss: 0.45890, Val Acc.: 85.20% | Epoch 46/75 | Train Loss: 0.37164, Train Acc.: 88.47% | Val Loss: 0.47161, Val Acc.: 84.30% | Epoch 47/75 | Train Loss: 0.37006, Train Acc.: 87.43% | Val Loss: 0.50671, Val Acc.: 83.44% | Epoch 48/75 | Train Loss: 0.35853, Train Acc.: 88.76% | Val Loss: 0.45435, Val Acc.: 84.78% | Epoch 49/75 | Train Loss: 0.36432, Train Acc.: 88.21% | Val Loss: 0.48283, Val Acc.: 84.02% | Epoch 50/75 | Train Loss: 0.36153, Train Acc.: 89.09% | Val Loss: 0.46196, Val Acc.: 84.78% | Epoch 51/75 | Train Loss: 0.35555, Train Acc.: 87.92% | Val Loss: 0.48926, Val Acc.: 83.24% | Epoch 52/75 | Train Loss: 0.35145, Train Acc.: 88.71% | Val Loss: 0.45024, Val Acc.: 84.84% | Epoch 53/75 | Train Loss: 0.35009, Train Acc.: 88.83% | Val Loss: 0.46385, Val Acc.: 84.44% | Epoch 54/75 | Train Loss: 0.34292, Train Acc.: 88.58% | Val Loss: 0.46576, Val Acc.: 84.28% | Epoch 55/75 | Train Loss: 0.34292, Train Acc.: 89.30% | Val Loss: 0.44746, Val Acc.: 85.26% | Epoch 56/75 | Train Loss: 0.35040, Train Acc.: 88.88% | Val Loss: 0.46846, Val Acc.: 83.84% | Epoch 57/75 | Train Loss: 0.34394, Train Acc.: 89.60% | Val Loss: 0.45665, Val Acc.: 84.68% | Epoch 58/75 | Train Loss: 0.33876, Train Acc.: 89.37% | Val Loss: 0.45056, Val Acc.: 85.10% | Epoch 59/75 | Train Loss: 0.34050, Train Acc.: 88.93% | Val Loss: 0.46470, Val Acc.: 84.48% | Epoch 60/75 | Train Loss: 0.33713, Train Acc.: 89.94% | Val Loss: 0.42983, Val Acc.: 85.12% | Epoch 61/75 | Train Loss: 0.33775, Train Acc.: 88.92% | Val Loss: 0.44717, Val Acc.: 84.88% | Epoch 62/75 | Train Loss: 0.32918, Train Acc.: 89.06% | Val Loss: 0.45263, Val Acc.: 84.08% | Epoch 63/75 | Train Loss: 0.32598, Train Acc.: 88.32% | Val Loss: 0.48561, Val Acc.: 83.90% | Epoch 64/75 | Train Loss: 0.32923, Train Acc.: 87.26% | Val Loss: 0.51722, Val Acc.: 82.96% | Epoch 65/75 | Train Loss: 0.32973, Train Acc.: 89.19% | Val Loss: 0.45999, Val Acc.: 84.92% | Epoch 66/75 | Train Loss: 0.32136, Train Acc.: 89.53% | Val Loss: 0.44131, Val Acc.: 85.42% | Epoch 67/75 | Train Loss: 0.32011, Train Acc.: 89.87% | Val Loss: 0.43502, Val Acc.: 85.46% | Epoch 68/75 | Train Loss: 0.32015, Train Acc.: 90.11% | Val Loss: 0.41001, Val Acc.: 85.72% | Epoch 69/75 | Train Loss: 0.31873, Train Acc.: 90.00% | Val Loss: 0.42961, Val Acc.: 85.10% | Epoch 70/75 | Train Loss: 0.31427, Train Acc.: 90.70% | Val Loss: 0.40964, Val Acc.: 86.04% | Epoch 71/75 | Train Loss: 0.31744, Train Acc.: 89.07% | Val Loss: 0.46226, Val Acc.: 84.74% | Epoch 72/75 | Train Loss: 0.31472, Train Acc.: 90.45% | Val Loss: 0.42773, Val Acc.: 85.56% | Epoch 73/75 | Train Loss: 0.31404, Train Acc.: 90.44% | Val Loss: 0.42573, Val Acc.: 86.12% | Epoch 74/75 | Train Loss: 0.30620, Train Acc.: 90.38% | Val Loss: 0.42480, Val Acc.: 85.70% | Epoch 75/75 | Train Loss: 0.31008, Train Acc.: 90.40% | Val Loss: 0.42157, Val Acc.: 85.60% | Total Training Time: 57.29 min Test Accuracy: 86.35%
In [16]:
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()
In [17]:
plot_history(history)
In [18]:
class_dict = {0: 'plane',
1: 'car',
2: 'bird',
3: 'cat',
4: 'deer',
5: 'dog',
6: 'frog',
7: 'horse',
8: 'ship',
9: 'truck'}
In [19]:
from sklearn.metrics import confusion_matrix
import seaborn as sns
def _compute_confusion_matrix(model, data_loader, device):
"""
Computes a confusion matrix for the given model and data.
Uses scikit-learn for robustness and speed.
"""
model.eval()
all_targets = []
all_predictions = []
with torch.no_grad():
for features, targets in data_loader:
features = features.to(device)
logits = model(features)
_, predicted_labels = torch.max(logits, 1)
# Move to CPU and convert to list/numpy immediately
all_targets.extend(targets.view(-1).cpu().numpy())
all_predictions.extend(predicted_labels.view(-1).cpu().numpy())
return confusion_matrix(all_targets, all_predictions)
def _plot_confusion_matrix(mat, class_names):
"""
Visualizes the confusion matrix using Seaborn for professional reporting.
"""
plt.figure(figsize=(10, 8))
sns.heatmap(mat, annot=True, fmt='d', cmap='Blues',
xticklabels=class_names, yticklabels=class_names)
plt.xlabel('Predicted Label')
plt.ylabel('True Label')
plt.title('Confusion Matrix: MobileNetV2 on CIFAR-10')
plt.show()
In [20]:
_mat = _compute_confusion_matrix(model=model, data_loader=test_loader, device=device)
_plot_confusion_matrix(_mat, class_names=class_dict.values())
In [21]:
%watermark -iv
matplotlib : 3.10.0 numpy : 2.0.2 pandas : 2.2.2 seaborn : 0.13.2 sklearn : 1.6.1 torch : 2.9.0+cu126 torchinfo : 1.8.0 torchvision: 0.24.0+cu126
In [21]: