EfficientNet V2 For Tensorflow2

This resource is using open-source code maintained in github (see the quick-start-guide section) and available for download from NGC

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.

Model architecture

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.

Default configuration

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

Feature support matrix

Features

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 training

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.

Enabling mixed precision

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)

Enabling TF32

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.