Training CIFAR-10 with TensorFlow2(TF2)

TensorFlow2 Classification Model Zoo. I’m playing with TensorFlow2 on the CIFAR-10 dataset.

Architectures

• LeNet
• AlexNet
• VGG [11, 13, 16, 19]
• ResNet [18, 34, 50, 101, 152]
• DenseNet [121, 169, 201]
• PreAct-ResNet [18, 34, 50, 101, 152]
• SENet
• SE-ResNet [18, 34, 50, 101, 152]
• SE-PreAct-ResNet [18, 34, 50, 101, 152]
• MobileNet
• MobileNetV2

Prerequisites

• Python 3.8+
• TensorFlow 2.4.0+

Training

Start training with:

python train.py --model resnet18

You can manually resume the training with:

python train.py --model resnet18 --resume

Testing

python test.py --model resnet18

Accuracy

Note

All abovementioned models are available. To specify the model, please use the model name without the hyphen. For instance, to train with SE-PreAct-ResNet18, you can run the following script:

python train.py --model sepreactresnet18

If you suffer from loss=nan issue, you can circumvent it by using a smaller learning rate, i.e.

python train.py --model sepreactresnet18 --lr 5e-2

GitHub

View Github