EfficientNet
Learn to train an EfficientNet image classification model.
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.
Get hands-on with 1200+ tech skills courses.