PyTorch Lightning CIFAR10

About this fork

Modified version of the PyTorch CIFAR project to exploit the PyTorch Lightning library.

In addition:

• Improvements main.py script, allowing you to train one or more models in a single command.
• Used PyTorch Lightning data module for the dataset (from the lightning-bolts package).
• Optimizer changed to use OneCycleLR scheduler.
• black formatter applied to all files.
• Added more consistent config of VGG and ShuffleNetV2 models.
• Added Tensorboard logging, and a JSON file final_metrics.json that saves final accuracy.

Library initially developed while at University of Glasgow's gicLAB.

Why is this useful?

CIFAR-10 is a small image classification dataset, which can be useful for validating research ideas (since models are smaller, and cheaper to train). It has two key advantages over the MNIST database in that the problem is a bit harder, and the images are non-greyscale. However, when doing research, you want to be able to have access to as much functionality as possible, without having to write a lot of boilerplate code. That is the design philosophy of the PyTorch Lightning library, hence why combining the two together makes sense. Note that few (if any) of these model architectures can be considered "official implementations", since the architectures have to be changed slightly to support the different data sizes. However, I have seen these models used in research, you need merely say that the models are defined for CIFAR10 (and ideally cite this repo, see the righthand side of the GitHub page!).

Contributing

PRs especially appreciated for features like:

• fixing ShuffleNetV1 (Currently having an issue with ShuffleNetV1, where initializing the model fails).
• expose more PyTorch Lightning features.
• Better logging systems.
• Systematic testing.
• CI/CD with Travis.

Prerequisites

• Python 3.6+
• PyTorch 1.7+
• lightning-bolts>=0.5.0
• pytorch-lightning>=1.6.2

Training

# Start training with:
python main.py --model_arch [your model, e.g. `mobilenetv2`, `resnet18 resnet50 vgg16`, or `all` for all models]

Accuracy

Models compared against the table reported in the original repo. Some models were not reported in the original repo, and are represented with a -