EfficientNet

Supported variants

The PyTorch Image Model supports the following EfficientNet variants:

Base model	Resolution
EfficientNetB0	224
EfficientNetB1	240
EfficientNetB2	260
EfficientNetB3	300
EfficientNetB4	380

EfficientNetB0 is a good start for most use cases. Other variants require higher GPU memory since the base image resolution is higher.

Training a new EfficientNet model

We can utilize the same training script to train a new EfficientNet image classification model. Run the following command in the terminal:

python train.py /app/dataset2 --model efficientnet_b0 --num-classes 4

Other variants

To train a different variant of EfficientNet, simply change the input for the model argument. For example, the following command will train a new EfficientNetB1 model:

python train.py /app/dataset2 --model efficientnet_b1 --num-classes 4

Custom epochs

By default, the number of epochs to train is 300. We can configure it by using the epochs argument. For example, the following command sets the training epochs to 500:

python train.py /app/dataset2 --model efficientnet_b0 --num-classes 4 --epochs 500

The AutoAugment augmentation

We can enable auto-augmentation by setting the aa flag as follows:

python train.py /app/dataset2 --model efficientnet_b0 --num-classes 4 --aa 'v0'

The Mixup or Cutmix augmentations

To enable Mixup augmentations, set the mixup argument as follows:

python train.py /app/dataset2 --model efficientnet_b0 --num-classes 4 --mixup 0.5

Set the cutmix argument for Cutmix augmentation:

python train.py /app/dataset2 --model efficientnet_b0 --num-classes 4 --cutmix 0.5

We can enable both augmentations with the following command:

python train.py /app/dataset2 --model efficientnet_b0 --num-classes 4 --mixup 0.5 --cutmix 0.5 --mixup-switch-prob 0.3

Note: Configure the mixup-switch-prob to switch between both augmentations randomly.

The Augmix augmentation

We can use the following command for Augmix augmentation:

python train.py /app/dataset2 --model efficientnet_b0 --num-classes 4 --aug-splits 3 --jsd

Most state-of-the-art models use augmentations as part of the training process. Augmentations help to improve the performance of our model.