Easy-to-run model benchmark on CIFAR10.
With this repository you can:
- train VGG[1], ResNet[2]
- manage training conditions using OmegaConf
- plot the results on tensorboard
- build environment using docker
So far you cannot:
- train models on ImageNet
- train models using multiple GPUs
- load tensorflow weights
- Updated docker and config environments (Dec. 2023)
- EfficientNet models are no longer supported (Dec. 2023)
RandomCrop and LRFlip are used for data augmentation.
Input size is 32x32.
| VGG16: 93.3 % | ResNet18: 94.3% |
|---|---|
|
|
- Python 3.6+
- PyTorch 1.0+
- (optional) tensorboardX
Install docker-compose and nvidia-container-runtime beforehand.
$ docker-compose build --build-arg UID="`id -u`" dev
$ docker-compose run devYou can select a model to train by specifying a config file.
$ python train.py --help
usage: train.py [-h] [--config CONFIG] [--tfboard TFBOARD]
[--checkpoint_dir CHECKPOINT_DIR] [--resume RESUME]
optional arguments:
-h, --help show this help message and exit
--config CONFIG path to config file
--tfboard TFBOARD tensorboard path for logging
--checkpoint_dir CHECKPOINT_DIR
directory where checkpoint files are saved
--resume RESUME checkpoint file pathexample:
$ python train.py --config configs/vgg16.yaml --tfboard out(TBD)
[1] K. Simonyan and A. Zisserman, "Very Deep Convolutional Networks for Large-Scale Image Recognition" paper
[2] K. He, X. Zhang, S. Ren, J. Sun, "Deep Residual Learning for Image Recognition" paper
[3] M. Tan and Q. V. Le, "EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks" paper
[4] pytorch-cifar