DEV Community

Peter Bao
Peter Bao

Posted on • Updated on

My Pytorch Training Pipeline in Image Classification

On this project, I create a custom dataset of 5 male models and conduct a full Pytorch training pipeline. I use a pretrained model and transfer learning, as well as do hyper-parameter search to help increase the accuracy.

ย 

โœ๏ธ Analysis and Evaluation:

  • Full Pytorch Training Pipeline on Image Classification part1

  • Full Pytorch Training Pipeline on Image Classification part2

๐Ÿ“บ My implementations are based off Aladdin Persson and Python Engineer.

ย 

My repository contains:

  • A training script (using pretrained vgg16 and transfer learning)
  • A script for Hyper-parameter Search
  • A script for loading the model for either resumed training or inference
  • A trained model ๐Ÿ˜ฌ
  • Some helper functions
  • A dataset

ย 

๐Ÿ“ฐ Dataset Structure

train/val
|___chau_minh_chi
       |___chau_minh_chi_01.jpg
          |___chau_minh_chi_02.jpg
          ...
|___keita_machida
       |___keita_machida_01.jpg
          |___keita_machida_02.jpg
          ...
Enter fullscreen mode Exit fullscreen mode

ย 

Small batch visualization

small

ย 

Large batch visualization

large

ย 

Visualize the number of classes

class

ย 

๐Ÿ˜‘ Hyper-parameter Search

Hyper-parameter search

ย 

๐Ÿ’ฅ Training Epochs

==> Saving new best
Epoch 1/25
Step 34/34, train Loss = 1.84,  train Acc = 0.29
Step 20/20, val loss = 1.58,  val acc = 0.25
Time spent for this epoch -----> 0m 32s

==> Saving new best
Epoch 2/25
Step 34/34, train Loss = 0.91,  train Acc = 0.63
Step 20/20, val loss = 1.53,  val acc = 0.43
Time spent for this epoch -----> 0m 13s

==> Validation accuracy did not improve.
Epoch 3/25
Step 34/34, train Loss = 0.61,  train Acc = 0.82
Step 20/20, val loss = 1.67,  val acc = 0.27
Time spent for this epoch -----> 0m 8s
Enter fullscreen mode Exit fullscreen mode

ย 

๐Ÿ˜Œ Visualize Loss and Accuracy

Loss and accuracy

ย 

๐Ÿ˜… Accuracy of Each Class

Test Acc
Got 13/30 correct samples over 43.33%
Accuracy of timmy_xu: 33.33%
Accuracy of corbyn_besson: 62.50%
Accuracy of keita_machida: 16.67%
Accuracy of wang_kai: 30.00%
Accuracy of chau_minh_chi: 100.00%
Enter fullscreen mode Exit fullscreen mode

ย 

๐Ÿ˜ž Classification Report and Confusion Matrix Heatmap

cm

ย 

๐Ÿ˜ Predict a single image

single

ย 

๐Ÿ˜ Prediction In Batch

In batch

Discussion (0)