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

Scripts and sample code

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

Parameters

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_v1/b0_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_v1/b0_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_v1/b0_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 want to change the values, override them via the command line.

Command-line options

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

Getting the data

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-*

Training process

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

By default, bash script scripts/{B0,B4}/training/{AMP,FP32,TF32}/convergence_8x{A100-80G,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/{B0,B4}/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

Multi-node runs can be launched on a Pyxis/enroot Slurm cluster (refer to Requirements) with the run_{B0,B4}_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_B0_multinode.sub
PARTITION=<partition_name> sbatch N 4 --ntasks-per-node=8 run_B4_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_{B0,B4}_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.