BUILD CIFAR-10 CLASSIFIER USING VARIOUS FRAMEWORKS

In this post, I share my collection of notebooks demonstrating how to build a CIFAR-10 classifier in various deep learning frameworks. The objectives are to show how to:

  • create a simple classifier using CNNs
  • track experiments
  • use hyperparameter tuning frameworks

I started working on this post to update my old notebooks on training a classifier on GPU vs TPU. In 2018, When I first wrote about GPU vs TPU, I wanted to find out the complexity involved in converting the code to switch from one accelerator (GPU) to another (TPUs). And whether it provided any benefit out of the box. A lot has changed since then.

TensorFlow 2.0 has made things a lot simpler. The eager mode is gentle on my brain, the Keras API, as always, is fun to work with. The introduction of tf.data API makes the construction of input pipelines easy. The features such as Autotune, cache, and prefetch take care of optimizing the pipeline. The tf.distribute.Strategy makes it simpler to switch between the accelerators (GPU, TPU).

This time around I decided to cover PyTorch, PyTorch Lightning, and JAX as well. While I do have some experience working with PyTorch and Lightning, JAX is mainly there because I wanted a reason to make something in JAX 😀.

Each card gives you some information about the notebook, training time, train and test accuracy, etc. I would advise you not to pay too much attention to the accuracy metrics because there is a slight difference in some notebooks' augmentation pipeline. Also, It is not my intention to perform any comparison between the frameworks. They all work great and may have pros and cons.

Update: Nov 3rd, 2020

My primary workstation, the one with GTX 1080TI in the cards below, is dead. I cannot continue with the following planned notebooks for now:

  • Optuna: PyTorch & PyTorch Lightning on GTX 1080TI
  • Ray Tune: PyTorch Lightning on GTX 1080TI
  • JAX * on GTX 1080TI

TENSORFLOW
TensorFlow Notebook

GTX 1080TI

Train Acc
73.95
Test Acc
67.83
Epochs
50
Training Time
4m 38s

Last Updated: 13 Aug 2020

GPUCOMET
TensorFlow + TPU Strategy

V2 8 Cores

Train Acc
78.29
Test Acc
68.49
Epochs
50
Training Time
4m 14s

Last Updated: 20 Sep 2020

TPUW&B
Keras Tuner Basic

GTX 1080TI

Train Acc
0
Test Acc
57.7
Epochs
50
Training Time
1h

Last Updated: 12 Sep 2020

GPU
Keras Tuner V2

GTX 1080TI

Train Acc
94.09
Test Acc
68.67
Epochs
50
Training Time
1h

Last Updated: 19 Sep 2020

GPU
Ray Tune

GTX 1080TI

Train Acc
100
Test Acc
61.7
Epochs
50
Training Time
1h

Last Updated: 03 Oct 2020

GPU
Optuna

GTX 1080TI

Train Acc
88.77
Test Acc
79.19
Epochs
50
Training Time
16m 8s

Last Updated: 19 Oct 2020

GPUACCURACY

originally published 14 Nov 2020 and updated 20 Dec 2020