torch-cifar-10-cnn

This repository contains a small CIFAR-10 / CIFAR-100 ResNet-like CNN implementation in PyTorch. Additionally, ResNets pre-trained on ImageNet can be finetuned to CIFAR-10 / CIFAR-100.

Some small pretrained models are available in the pretrained_models directory.

See https://github.com/menzHSE/mlx-cifar-10-cnn for (more or less) the same model being trained using MLX.

Requirements

torch
torchvision
torchinfo
numpy
Pillow

See requirements.txt

Usage

$ python train.py -h
usage: train.py [-h] [--cpu] [--seed SEED] [--batchsize BATCHSIZE] [--epochs EPOCHS] [--lr LR]
                [--dataset {CIFAR-10,CIFAR-100}] [--finetune {resnet18,resnet34,resnet50}]

Train a simple CNN on CIFAR-10 / CIFAR_100 with PyTorch.

options:
  -h, --help            show this help message and exit
  --cpu                 Use CPU instead of GPU (cuda/mps) acceleration
  --seed SEED           Random seed
  --batchsize BATCHSIZE
                        Batch size for training
  --epochs EPOCHS       Number of training epochs
  --lr LR               Learning rate
  --dataset {CIFAR-10,CIFAR-100}
                        Select the dataset to use (CIFAR-10 or CIFAR-100)
  --finetune {resnet18,resnet34,resnet50}
                        Select the model for fine-tuning (resnet18, resnet34, resnet50), omit for training
                        from scratch

$ python test.py -h
usage: test.py [-h] [--cpu] --model MODEL [--dataset {CIFAR-10,CIFAR-100}]
               [--finetune {resnet18,resnet34,resnet50}]

Test a simple CNN on CIFAR-10 / CIFAR-100 with PyTorch.

options:
  -h, --help            show this help message and exit
  --cpu                 Use CPU instead of GPU (cuda/mps) acceleration
  --model MODEL         Model filename *.pth
  --dataset {CIFAR-10,CIFAR-100}
                        Select the dataset to use (CIFAR-10 or CIFAR-100)
  --finetune {resnet18,resnet34,resnet50}
                        Select the model that has been fine-tuned(resnet18, resnet34, resnet50), omit for
                        testing the original model

Examples

CIFAR-10

Train on CIFAR-10

python train.py

This uses an 80GB NVIDIA A100.

$  python train.py --batchsize=512  --dataset CIFAR-10
Using device: cuda
NVIDIA A100 80GB PCIe
Options:
  Device: GPU
  Seed: 0
  Batch size: 512
  Number of epochs: 30
  Learning rate: 0.0003
  Dataset: CIFAR-10
Files already downloaded and verified
Files already downloaded and verified
==========================================================================================
Layer (type (var_name))                  Output Shape              Param #
==========================================================================================
CNN (CNN)                                [1, 10]                   --
├─Conv2d (conv1)                         [1, 32, 32, 32]           896
├─BatchNorm2d (bn1)                      [1, 32, 32, 32]           64
├─Conv2d (skip2)                         [1, 32, 32, 32]           1,056
├─Conv2d (conv2)                         [1, 32, 32, 32]           9,248
├─BatchNorm2d (bn2)                      [1, 32, 32, 32]           64
├─MaxPool2d (pool)                       [1, 32, 16, 16]           --
├─Dropout (drop)                         [1, 32, 16, 16]           --
├─Conv2d (conv3)                         [1, 64, 16, 16]           18,496
├─BatchNorm2d (bn3)                      [1, 64, 16, 16]           128
├─Conv2d (skip4)                         [1, 64, 16, 16]           4,160
├─Conv2d (conv4)                         [1, 64, 16, 16]           36,928
├─BatchNorm2d (bn4)                      [1, 64, 16, 16]           128
├─MaxPool2d (pool)                       [1, 64, 8, 8]             --
├─Dropout (drop)                         [1, 64, 8, 8]             --
├─Conv2d (conv5)                         [1, 128, 8, 8]            73,856
├─BatchNorm2d (bn5)                      [1, 128, 8, 8]            256
├─Conv2d (skip6)                         [1, 128, 8, 8]            16,512
├─Conv2d (conv6)                         [1, 128, 8, 8]            147,584
├─BatchNorm2d (bn6)                      [1, 128, 8, 8]            256
├─MaxPool2d (pool)                       [1, 128, 4, 4]            --
├─Dropout (drop)                         [1, 128, 4, 4]            --
├─Linear (fc1)                           [1, 128]                  262,272
├─Linear (fc2)                           [1, 10]                   1,290
==========================================================================================
Total params: 573,194
Trainable params: 573,194
Non-trainable params: 0
Total mult-adds (M): 42.22
==========================================================================================
Input size (MB): 0.01
Forward/backward pass size (MB): 2.29
Params size (MB): 2.29
Estimated Total Size (MB): 4.60
==========================================================================================

[Epoch   0] :  . done (98 batches)
[Epoch   0] : | time:  3.463s | trainLoss:  1.608 | trainAccuracy:  0.409 | valLoss:  1.547 | valAccuracy:  0.460 | throughput:  86253.256 img/s |
[Epoch   1] :  . done (98 batches)
[Epoch   1] : | time:  3.141s | trainLoss:  1.193 | trainAccuracy:  0.573 | valLoss:  1.258 | valAccuracy:  0.567 | throughput:  83937.610 img/s |
[Epoch   2] :  . done (98 batches)
[Epoch   2] : | time:  3.375s | trainLoss:  1.008 | trainAccuracy:  0.642 | valLoss:  0.969 | valAccuracy:  0.668 | throughput:  68932.503 img/s |
[Epoch   3] :  . done (98 batches)
[Epoch   3] : | time:  3.196s | trainLoss:  0.886 | trainAccuracy:  0.687 | valLoss:  0.954 | valAccuracy:  0.677 | throughput:  73139.131 img/s |
[Epoch   4] :  . done (98 batches)
[Epoch   4] : | time:  3.349s | trainLoss:  0.787 | trainAccuracy:  0.724 | valLoss:  0.924 | valAccuracy:  0.686 | throughput:  70358.619 img/s |
[Epoch   5] :  . done (98 batches)
[Epoch   5] : | time:  3.202s | trainLoss:  0.717 | trainAccuracy:  0.748 | valLoss:  0.840 | valAccuracy:  0.717 | throughput:  73180.432 img/s |
[Epoch   6] :  . done (98 batches)
[Epoch   6] : | time:  3.227s | trainLoss:  0.660 | trainAccuracy:  0.769 | valLoss:  0.699 | valAccuracy:  0.764 | throughput:  85500.814 img/s |
[Epoch   7] :  . done (98 batches)
[Epoch   7] : | time:  3.094s | trainLoss:  0.612 | trainAccuracy:  0.787 | valLoss:  0.662 | valAccuracy:  0.778 | throughput:  86023.888 img/s |
[Epoch   8] :  . done (98 batches)
[Epoch   8] : | time:  3.067s | trainLoss:  0.567 | trainAccuracy:  0.803 | valLoss:  0.630 | valAccuracy:  0.786 | throughput:  86963.361 img/s |
[Epoch   9] :  . done (98 batches)
[Epoch   9] : | time:  3.234s | trainLoss:  0.527 | trainAccuracy:  0.816 | valLoss:  0.606 | valAccuracy:  0.796 | throughput:  81349.960 img/s |
[Epoch  10] :  . done (98 batches)
[Epoch  10] : | time:  3.127s | trainLoss:  0.498 | trainAccuracy:  0.828 | valLoss:  0.604 | valAccuracy:  0.798 | throughput:  86073.199 img/s |
[Epoch  11] :  . done (98 batches)
[Epoch  11] : | time:  3.115s | trainLoss:  0.469 | trainAccuracy:  0.836 | valLoss:  0.584 | valAccuracy:  0.803 | throughput:  80480.693 img/s |
[Epoch  12] :  . done (98 batches)
[Epoch  12] : | time:  3.151s | trainLoss:  0.442 | trainAccuracy:  0.846 | valLoss:  0.577 | valAccuracy:  0.806 | throughput:  77756.301 img/s |
[Epoch  13] :  . done (98 batches)
[Epoch  13] : | time:  3.211s | trainLoss:  0.415 | trainAccuracy:  0.855 | valLoss:  0.564 | valAccuracy:  0.811 | throughput:  73234.967 img/s |
[Epoch  14] :  . done (98 batches)
[Epoch  14] : | time:  3.118s | trainLoss:  0.385 | trainAccuracy:  0.865 | valLoss:  0.549 | valAccuracy:  0.817 | throughput:  82200.094 img/s |
[Epoch  15] :  . done (98 batches)
[Epoch  15] : | time:  3.241s | trainLoss:  0.371 | trainAccuracy:  0.869 | valLoss:  0.543 | valAccuracy:  0.821 | throughput:  75561.320 img/s |
[Epoch  16] :  . done (98 batches)
[Epoch  16] : | time:  3.442s | trainLoss:  0.351 | trainAccuracy:  0.877 | valLoss:  0.552 | valAccuracy:  0.821 | throughput:  64772.159 img/s |
[Epoch  17] :  . done (98 batches)
[Epoch  17] : | time:  3.284s | trainLoss:  0.335 | trainAccuracy:  0.883 | valLoss:  0.542 | valAccuracy:  0.827 | throughput:  71640.955 img/s |
[Epoch  18] :  . done (98 batches)
[Epoch  18] : | time:  3.151s | trainLoss:  0.318 | trainAccuracy:  0.890 | valLoss:  0.556 | valAccuracy:  0.820 | throughput:  77383.137 img/s |
[Epoch  19] :  . done (98 batches)
[Epoch  19] : | time:  3.132s | trainLoss:  0.303 | trainAccuracy:  0.894 | valLoss:  0.531 | valAccuracy:  0.830 | throughput:  86589.410 img/s |
[Epoch  20] :  . done (98 batches)
[Epoch  20] : | time:  3.293s | trainLoss:  0.283 | trainAccuracy:  0.901 | valLoss:  0.536 | valAccuracy:  0.830 | throughput:  67708.022 img/s |
[Epoch  21] :  . done (98 batches)
[Epoch  21] : | time:  3.192s | trainLoss:  0.272 | trainAccuracy:  0.905 | valLoss:  0.551 | valAccuracy:  0.821 | throughput:  80829.243 img/s |
[Epoch  22] :  . done (98 batches)
[Epoch  22] : | time:  3.226s | trainLoss:  0.253 | trainAccuracy:  0.911 | valLoss:  0.555 | valAccuracy:  0.827 | throughput:  72563.493 img/s |
[Epoch  23] :  . done (98 batches)
[Epoch  23] : | time:  3.178s | trainLoss:  0.247 | trainAccuracy:  0.913 | valLoss:  0.558 | valAccuracy:  0.832 | throughput:  86425.289 img/s |
[Epoch  24] :  . done (98 batches)
[Epoch  24] : | time:  3.156s | trainLoss:  0.235 | trainAccuracy:  0.918 | valLoss:  0.565 | valAccuracy:  0.823 | throughput:  85690.960 img/s |
[Epoch  25] :  . done (98 batches)
[Epoch  25] : | time:  3.200s | trainLoss:  0.227 | trainAccuracy:  0.921 | valLoss:  0.545 | valAccuracy:  0.835 | throughput:  86722.386 img/s |
[Epoch  26] :  . done (98 batches)
[Epoch  26] : | time:  3.311s | trainLoss:  0.216 | trainAccuracy:  0.925 | valLoss:  0.579 | valAccuracy:  0.825 | throughput:  86752.040 img/s |
[Epoch  27] :  . done (98 batches)
[Epoch  27] : | time:  3.223s | trainLoss:  0.203 | trainAccuracy:  0.929 | valLoss:  0.523 | valAccuracy:  0.843 | throughput:  73632.212 img/s |
[Epoch  28] :  . done (98 batches)
[Epoch  28] : | time:  3.493s | trainLoss:  0.191 | trainAccuracy:  0.933 | valLoss:  0.585 | valAccuracy:  0.828 | throughput:  60934.829 img/s |
[Epoch  29] :  . done (98 batches)
[Epoch  29] : | time:  3.488s | trainLoss:  0.189 | trainAccuracy:  0.933 | valLoss:  0.557 | valAccuracy:  0.837 | throughput:  63697.672 img/s |
Training finished in 0:01:36.280061 hh:mm:ss.ms

Test on CIFAR-10

$ python test.py --model=models/model_CIFAR-10_029.pth
Using device: cuda
NVIDIA A100 80GB PCIe
Files already downloaded and verified
Files already downloaded and verified
Loaded model for CIFAR-10 from models/model_CIFAR-10_029.pth
Testing ....
Test Metrics (10000 test samples):
  - Accuracy: 0.837
Testing finished in 0:00:02.602382 hh:mm:ss.ms

CIFAR-100

Train on CIFAR-100

python train.py --dataset CIFAR-100