This repository provides scripts and recipes to train EfficientNet v2-S to achieve state-of-the-art accuracy. The content of the repository is maintained by NVIDIA and is tested against each NGC monthly released container to ensure consistent accuracy and performance over time.

• Model overview

  • Model architecture
  • Default configuration
  • Feature support matrix
    • Features
  • Mixed precision training
    • Enabling mixed precision
    • Enabling TF32

• Setup

  • Requirements

• Quick Start Guide

• Advanced

  • Scripts and sample code
  • Parameters
  • Command-line options
  • Getting the data
  • Training process
    • Multi-node
  • Inference process

• Performance

  • Benchmarking
  • Training performance benchmark
    • Inference performance benchmark
  • Results
    • Training results for EfficientNet v2-S
      • Training accuracy: NVIDIA DGX A100 (8x A100 80GB)
      • Training accuracy: NVIDIA DGX-1 (8x V100 32GB)
    • Training performance results for EfficientNet v2-S
      • Training performance: NVIDIA DGX A100 (8x A100 80GB)
      • Training performance: NVIDIA DGX-1 (8x V100 32GB)
    • Training EfficientNet v2-S at scale
      • 10x NVIDIA DGX-1 V100 (8x V100 32GB)
      • 10x NVIDIA DGX A100 (8x A100 80GB)
    • Inference performance results for EfficientNet v2-S
      • Inference performance: NVIDIA DGX A100 (1x A100 80GB)
      • Inference performance: NVIDIA DGX-1 (1x V100 32GB)

• Release notes

  • Changelog
  • Known issues

EfficientNet TensorFlow 2 is a family of image classification models, which achieve state-of-the-art accuracy, yet being an order-of-magnitude smaller and faster than previous models. Specifically, this readme covers model v2-S as suggested in EfficientNetV2: Smaller Models and Faster Training. NVIDIA's implementation of EfficientNet TensorFlow 2 is an optimized version of TensorFlow Model Garden implementation, leveraging mixed precision arithmetic on NVIDIA Volta, NVIDIA Turing, and the NVIDIA Ampere GPU architectures for faster training times while maintaining target accuracy.

The major differences between the papers' original implementations and this version of EfficientNet are as follows:

• Automatic mixed precision (AMP) training support
• Cosine LR decay for better accuracy
• Weight initialization using fan_out for better accuracy
• Multi-node training support using Horovod
• XLA enabled for better performance
• Gradient accumulation support
• Lightweight logging using dllogger

Other publicly available implementations of EfficientNet include:

• Tensorflow Model Garden
• Pytorch version
• Google's implementation for TPU EfficientNet v1
• Google's implementation for TPU EfficientNet v2

This model is trained with mixed precision Tensor Cores on NVIDIA Volta, NVIDIA Turing, and the NVIDIA Ampere GPU architectures. It provides a push-button solution to pretraining on a corpus of choice. As a result, researchers can get results 1.5--2x faster than training without Tensor Cores, while experiencing the benefits of mixed precision training. This model is tested against each NGC monthly released container to ensure consistent accuracy and performance over time.

EfficientNet v2 is developed based on AutoML and compound scaling, but with a particular emphasis on faster training. For this purpose, the authors have proposed 3 major changes compared to v1: 1) the objective function of AutoML is revised so that the number of flops is now substituted by training time, because FLOPs is not an accurate surrogate of the actual training time; 2) a multi-stage training is proposed where the early stages of training use low resolution images and weak regularization, but the subsequent stages use larger images and stronger regularization; 3) an additional block called fused MBConv is used in AutoML, which replaces the 1x1 depth-wise convolution of MBConv with a regular 3x3 convolution.

EfficientNet v2 base model is scaled up using a non-uniform compounding scheme, through which the depth and width of blocks are scaled depending on where they are located in the base architecture. With this approach, the authors have identified the base "small" model, EfficientNet v2-S, and then scaled it up to obtain EfficientNet v2-M,L,XL. Below is the detailed overview of EfficientNet v2-S, which is reproduced in this repository.

Here is the baseline EfficientNet v2-S structure.

The following features are supported by this implementation:

• General:

  • XLA support
  • Mixed precision support
  • Multi-GPU support using Horovod
  • Multi-node support using Horovod
  • Cosine LR Decay

• Inference:

  • Support for inference on a single image is included
  • Support for inference on a batch of images is included

Multi-GPU training with Horovod

Our model uses Horovod to implement efficient multi-GPU training with NCCL. For details, refer to example sources in this repository or refer to the TensorFlow tutorial.

Multi-node training with Horovod

Our model also uses Horovod to implement efficient multi-node training.

Automatic Mixed Precision (AMP)

Computation graphs can be modified by TensorFlow on runtime to support mixed precision training. A detailed explanation of mixed precision can be found in Appendix.

Gradient Accumulation

Gradient Accumulation is supported through a custom train_step function. This feature is enabled only when grad_accum_steps is greater than 1.

Stage-wise Training

Stage-wise training was proposed for EfficientNet v2 to further accelerate convergence. In this scheme, the early stages use low resolution images and weak regularization, but the subsequent stages use larger images and stronger regularization. This feature is activated when --n_stages is greater than 1. The current codebase allows the user to linearly schedule the following factors in the various stages of training:

Note that if --n_stages is set to 1, then the above hyper-parameters beginning with base_ will have no effect.

Mixed precision is the combined use of different numerical precisions in a computational method. Mixed precision training offers significant computational speedup by performing operations in half-precision format while storing minimal information in single-precision to retain as much information as possible in critical parts of the network. Since the introduction of Tensor Cores in NVIDIA Volta, and following with both the NVIDIA Turing and NVIDIA Ampere Architectures, significant training speedups are experienced by switching to mixed precision -- up to 3x overall speedup on the most arithmetically intense model architectures. Using mixed precision training previously required two steps:

1. Porting the model to use the FP16 data type where appropriate.
2. Adding loss scaling to preserve small gradient values.

This can now be achieved using Automatic Mixed Precision (AMP) for TensorFlow to enable the full mixed precision methodology in your existing TensorFlow model code. AMP enables mixed precision training on NVIDIA Volta, NVIDIA Turing, and NVIDIA Ampere GPU architectures automatically. The TensorFlow framework code makes all necessary model changes internally.

In TF-AMP, the computational graph is optimized to use as few casts as necessary and maximize the use of FP16, and the loss scaling is automatically applied inside of supported optimizers. AMP can be configured to work with the existing tf.contrib loss scaling manager by disabling the AMP scaling with a single environment variable to perform only the automatic mixed-precision optimization. It accomplishes this by automatically rewriting all computation graphs with the necessary operations to enable mixed precision training and automatic loss scaling.

For information about:

• How to train using mixed precision, refer to the Mixed Precision Training paper and Training With Mixed Precision documentation.
• Techniques used for mixed precision training, refer to the Mixed-Precision Training of Deep Neural Networks blog.
• How to access and enable AMP for TensorFlow, refer to Using TF-AMP from the TensorFlow User Guide.

Mixed precision is enabled in TensorFlow by using the Automatic Mixed Precision (TF-AMP) extension which casts variables to half-precision upon retrieval, while storing variables in single-precision format. Furthermore, to preserve small gradient magnitudes in backpropagation, a loss scaling step must be included when applying gradients. In TensorFlow, loss scaling can be applied statically by using simple multiplication of loss by a constant value or automatically, by TF-AMP. Automatic mixed precision makes all the adjustments internally in TensorFlow, providing two benefits over manual operations. First, programmers need not modify network model code, reducing development and maintenance effort. Second, using AMP maintains forward and backward compatibility with all the APIs for defining and running TensorFlow models.

To enable mixed precision, you can simply add the --use_amp to the command-line used to run the model. This will enable the following code:

if params.use_amp:
    policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16', loss_scale='dynamic')
    tf.keras.mixed_precision.experimental.set_policy(policy)

TensorFloat-32 (TF32) is the new math mode in NVIDIA A100 GPUs for handling the matrix math also called tensor operations. TF32 running on Tensor Cores in A100 GPUs can provide up to 10x speedups compared to single-precision floating-point math (FP32) on NVIDIA Volta GPUs.

TF32 Tensor Cores can speed up networks using FP32, typically with no loss of accuracy. It is more robust than FP16 for models which require a high dynamic range for weights or activations.

For more information, refer to the TensorFloat-32 in the A100 GPU Accelerates AI Training, HPC up to 20x blog post.

TF32 is supported in the NVIDIA Ampere GPU architecture and is enabled by default.

The following section lists the requirements that you need to meet in order to start training the EfficientNet model.

This repository contains a Dockerfile that extends the TensorFlow NGC container and encapsulates some dependencies. Aside from these dependencies, ensure you have the following components:

• NVIDIA Docker
• [TensorFlow 21.09-py3] NGC container or later
• Supported GPUs:
  • NVIDIA Volta architecture
  • NVIDIA Turing architecture
  • NVIDIA Ampere architecture

For more information about how to get started with NGC containers, refer to the following sections from the NVIDIA GPU Cloud Documentation and the Deep Learning Documentation:

• Getting Started Using NVIDIA GPU Cloud
• Accessing And Pulling From The NGC Container Registry
• Running TensorFlow

As an alternative to the use of the Tensorflow2 NGC container, to set up the required environment or create your own container, refer to the versioned NVIDIA Container Support Matrix.

For multi-node, the sample provided in this repository requires Enroot and Pyxis set up on a SLURM cluster.

To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the EfficientNet model on the ImageNet dataset. For the specifics concerning training and inference, refer to the Advanced section.

1. Clone the repository.

   git clone https://github.com/NVIDIA/DeepLearningExamples.git
   
   cd DeepLearningExamples/TensorFlow2/Classification/ConvNets/efficientnet
   

2. Download and prepare the dataset. Runner.py supports ImageNet with TensorFlow Datasets (TFDS). Refer to the TFDS ImageNet readme for manual download instructions.

3. Build EfficientNet on top of the NGC container. bash ./scripts/docker/build.sh YOUR_DESIRED_CONTAINER_NAME

4. Start an interactive session in the NGC container to run training/inference. Ensure that launch.sh has the correct path to ImageNet on your machine and that this path is mounted onto the /data directory, because this is where training and evaluation scripts search for data.

       `bash ./scripts/docker/launch.sh YOUR_DESIRED_CONTAINER_NAME`
   

5. Start training.

   To run training for a standard configuration, under the container default entry point /workspace, run one of the scripts in the ./efficinetnet_v2/S/training/{AMP,TF32,FP32}/convergence_8x{A100-80G, V100-16G, V100-32G}.sh. For example:

   bash ./efficinetnet_v2/S/training/AMP/convergence_8xA100-80G.sh

6. Start validation/evaluation.

   To run validation/evaluation for a standard configuration, under the container default entry point /workspace, run one of the scripts in the ./efficinetnet_v2/S/evaluation/evaluation_{AMP,FP32,TF32}_8x{A100-80G,V100-16G,V100-32G}.sh. The evaluation script is configured to use the checkpoint specified in the checkpoint file for evaluation. The specified checkpoint will be read from the location passed by `--model_dir'.For example:

   bash ./efficinetnet_v2/S/evaluation/evaluation_AMP_A100-80G.sh

7. Start inference/predictions.

   To run inference for a standard configuration, under the container default entry point /workspace, run one of the scripts in the ./efficinetnet_v2/S/inference/inference_{AMP,FP32,TF32}.sh. Ensure your JPEG images used to run inference on are mounted in the /infer_data directory with this folder structure :

   infer_data
   |   ├── images
   |   |   ├── image1.JPEG
   |   |   ├── image2.JPEG
   

   For example: bash ./efficinetnet_v2/S/inference/inference_AMP.sh

Now that you have your model trained and evaluated, you can choose to compare your training results with our Training accuracy results. You can also choose to benchmark your performance to Training performance benchmark, or Inference performance benchmark. Following the steps in these sections will ensure that you achieve the same accuracy and performance results as stated in the Results section.

The following sections provide greater details of the dataset, running training and inference, and the training results.

The repository is structured as follows:

• scripts/ - shell scripts to build and launch EfficientNet container on top of NGC container,
• efficientnet_{v1,v2} scripts to launch training, evaluation and inference
• model/ - building blocks and EfficientNet model definitions
• runtime/ - detailed procedure for each running mode
• utils/ - support util functions for learning rates, optimizers, etc.
• dataloader/ provides data pipeline utils
• config/ contains model definitions

The hyper parameters can be grouped into model-specific hyperparameters (e.g., #layers ) and general hyperparameters (e.g., #training epochs).

The model-specific hyperparameters are to be defined in a python module, which must be passed in the command line via --cfg ( python main.py --cfg config/efficientnet_v2/s_cfg.py). To override model-specific hyperparameters, you can use a comma separated list of k=v pairs (e.g., python main.py --cfg config/efficientnet_v2/s_cfg.py --mparams=bn_momentum=0.9,dropout=0.5).

The general hyperparameters and their default values can be found in utils/cmdline_helper.py. The user can override these hyperparameters in the command line (e.g., python main.py --cfg config/efficientnet_v2/s_cfg.py --data_dir xx --train_batch_size 128). Here is a list of important hyperparameters:

• --mode (train_and_eval,train,eval,prediction) - the default is train_and_eval.

• --use_amp Set to True to enable AMP

• --use_xla Set to True to enable XLA

• --model_dir The folder where model checkpoints are saved (the default is /workspace/output)

• --data_dir The folder where data resides (the default is /data/)

• --log_steps The interval of steps between logging of batch level stats.

• --augmenter_name Type of data augmentation

• --raug_num_layers Number of layers used in the random data augmentation scheme

• --raug_magnitude Strength of transformations applied in the random data augmentation scheme

• --cutmix_alpha Cutmix parameter used in the last stage of training.

• --mixup_alpha Mixup parameter used in the last stage of training.

• --defer_img_mixing Move image mixing ops from the data loader to the model/GPU (faster training)

• --eval_img_size Size of images used for evaluation

• --eval_batch_size The evaluation batch size per GPU

• --n_stages Number of stages used for stage-wise training

• --train_batch_size The training batch size per GPU

• --train_img_size Size of images used in the last stage of training

• --base_img_size Size of images used in the first stage of training

• --max_epochs The number of training epochs

• --warmup_epochs The number of epochs of warmup

• --moving_average_decay The decay weight used for EMA

• --lr_init The learning rate for a batch size of 128, effective learning rate will be automatically scaled according to the global training batch size: lr=lr_init * global_BS/128 where global_BS=train_batch_size*n_GPUs

• --lr_decay Learning rate decay policy

• --weight_decay Weight decay coefficient

• --save_checkpoint_freq Number of epochs to save checkpoints

NOTE: Avoid changing the default values of the general hyperparameters provided in utils/cmdline_helper.py. The reason is that some other models supported by this repository may rely on such default values. If you wish to change the values, override them via the command line.

To display the full list of available options and their descriptions, use the -h or --help command-line option, for example: python main.py --help

Refer to the TFDS ImageNet readme for manual download instructions. To train on the ImageNet dataset, pass $path_to_ImageNet_tfrecords to $data_dir in the command-line.

Name the TFRecords in the following scheme:

• Training images - /data/train-*
• Validation images - /data/validation-*

The training process can start from scratch, or resume from a checkpoint.

By default, bash script scripts/S/training/{AMP,FP32,TF32}/convergence_8x{A100-80G,V100-16G,V100-32G}.sh will start the training process with the following settings.

• Use 8 GPUs by Horovod
• Has XLA enabled
• Saves checkpoints after every 10 epochs to /workspace/output/ folder
• AMP or FP32 or TF32 based on the folder scripts/S/training/{AMP, FP32, TF32}

The training starts from scratch if --model_dir has no checkpoints in it. To resume from a checkpoint, place the checkpoint into --model_dir and make sure the checkpoint file points to it.

Multi-node runs can be launched on a Pyxis/enroot Slurm cluster (refer to Requirements) with the run_S_multinode.sub script with the following command for a 4-node NVIDIA DGX A100 example:

PARTITION=<partition_name> sbatch N 4 --ntasks-per-node=8 run_S_multinode.sub

Checkpoints will be saved after --save_checkpoint_freq epochs at checkpointdir. The latest checkpoint will be automatically picked up to resume training in case it needs to be resumed. Cluster partition name has to be provided <partition_name>.

Note that the run_S_multinode.sub script is a starting point that has to be adapted depending on the environment. In particular, pay attention to the variables such as --container-image, which handles the container image to train, and --datadir, which handles the location of the ImageNet data.

Refer to the scripts to find the full list of variables to adjust for your system.

Validation can be done either during training (when --mode train_and_eval is used) or in a post-training setting (--mode eval) on a checkpointed model. The evaluation script expects data in the tfrecord format.

bash ./scripts/S/evaluation/evaluation_{AMP,FP32,TF32}_{A100-80G,V100-16G,V100-32G}.sh

Metrics gathered through this process are listed below:

- eval_loss
- eval_accuracy_top_1
- eval_accuracy_top_5
- avg_exp_per_second_eval
- avg_exp_per_second_eval_per_GPU
- avg_time_per_exp_eval : Average Latency
- latency_90pct : 90% Latency
- latency_95pct : 95% Latency
- latency_99pct : 99% Latency

The scripts used for inference expect the inference data in the following directory structure:

```
infer_data
|   ├── images
|   |   ├── image1.JPEG
|   |   ├── image2.JPEG
```

Run: bash ./scripts/S/inference/inference_{AMP,FP32,TF32}.sh

The performance measurements in this document were conducted at the time of publication and may not reflect the performance achieved from NVIDIA’s latest software release. For the most up-to-date performance measurements, go to NVIDIA Data Center Deep Learning Product Performance.

The following section shows how to run benchmarks measuring the model performance in training and inference modes.

Training benchmark for EfficientNet v2-S was run on NVIDIA DGX A100 80GB and NVIDIA DGX-1 V100 32GB.

bash ./scripts/S/training/{AMP, FP32, TF32}/train_benchmark_8x{A100-80G, V100-32G}.sh

Inference benchmark for EfficientNet v2-S was run on NVIDIA DGX A100 80GB and NVIDIA DGX-1 V100 32GB.

The following sections provide details on how we achieved our performance and accuracy in training and inference.

Our results were obtained by running the training scripts in the tensorflow:21.09-tf2-py3 NGC container on multi-node NVIDIA DGX A100 (8x A100 80GB) GPUs. We evaluated the models using both the original and EMA weights and selected the higher accuracy to report.

Our results were obtained by running the training scripts in the tensorflow:21.09-tf2-py3 NGC container on multi-node NVIDIA DGX V100 (8x V100 32GB) GPUs. We evaluated the models using both the original and EMA weights and selected the higher accuracy to report.

EfficientNet v2-S uses images of increasing resolution during training. Since throughput changes depending on the image size, we have measured throughput based on the image size used in the last stage of training (300x300).

EfficientNet v2-S uses images of increasing resolution during training. Since throughput changes depending on the image size, we have measured throughput based on the image size used in the last stage of training (300x300).

We trained EfficientNet v2-S at scale using 10 DGX-1 machines each having 8x V100 32GB GPUs. We used the same set of hyperparameters and NGC container as before. Also, throughput numbers were measured in the last stage of training. The accuracy was selected as the better between that of the original weights and EMA weights.

We trained EfficientNet v2-S at scale using 10 DGX A100 machines each having 8x A100 80GB GPUs. This training setting has an effective batch size of 36800 (460x8x10), which requires advanced optimizers particularly designed for large-batch training. For this purpose, we used the nvLAMB optimizer with the following hyper parameters: lr_warmup_epochs=10, beta_1=0.9, beta_2=0.999, epsilon=0.000001, grad_global_clip_norm=1, lr_init=0.00005, weight_decay=0.00001. As before, we used tensorflow:21.09-tf2-py3 NGC container and measured throughput numbers in the last stage of training. The accuracy was selected as the better between that of the original weights and EMA weights.

Our results were obtained by running the inferencing benchmarking script in the tensorflow:21.09-tf2-py3 NGC container on the NVIDIA DGX A100 (1x A100 80GB) GPU.

FP16 Inference Latency

TF32 Inference Latency

Our results were obtained by running the inferencing benchmarking script in the tensorflow:21.09-tf2-py3 NGC container on the NVIDIA DGX V100 (1x V100 32GB) GPU.

FP16 Inference Latency

FP32 Inference Latency

February 2022

• Initial release

• EfficientNet v2 might run into OOM or might exhibit a significant drop in throughput at the onset of the last stage of training. The fix is to resume training from the latest checkpoint.