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

Scripts and sample code

For a non-standard configuration, run:

• For 1 GPU

  • FP32 python ./main.py --arch efficientnet-<version> --label-smoothing 0.1 <path to imagenet> python ./main.py --arch efficientnet-<version> --label-smoothing 0.1 --amp --static-loss-scale 256 <path to imagenet>

• For multiple GPUs

  • FP32 python ./multiproc.py --nproc_per_node 8 ./main.py --arch efficientnet-<version> --label-smoothing 0.1 <path to imagenet>
  • AMP python ./multiproc.py --nproc_per_node 8 ./main.py --arch efficientnet-<version> --label-smoothing 0.1 --amp --static-loss-scale 256 <path to imagenet>

Use python ./main.py -h to obtain the list of available options in the main.py script.

Command-line options

To see the full list of available options and their descriptions, use the -h or --help command-line option, for example:

python main.py -h

Dataset guidelines

To use your own dataset, divide it into directories. For example:

• Training images - train/<class id>/<image>
• Validation images - val/<class id>/<image>

If your dataset has a number of classes different than 1000, you need to pass the --num_classes N flag to the training script.

Training process

All the results of the training will be stored in the directory specified with --workspace argument.

The script will store:

• the most recent checkpoint - checkpoint.pth.tar (unless --no-checkpoints flag is used).
• the checkpoint with the best validation accuracy - model_best.pth.tar (unless --no-checkpoints flag is used).
• the JSON log - in the file specified with the --raport-file flag.

Metrics gathered through training:

• train.loss - training loss
• train.total_ips - training speed measured in images/second
• train.compute_ips - training speed measured in images/second, not counting data loading
• train.data_time - time spent on waiting on data
• train.compute_time - time spent in forward/backward pass

To restart training from the checkpoint use the --resume option.

To start training from pretrained weights (for example, downloaded from NGC) use the --pretrained-from-file option.

The difference between --resume and --pretrained-from-file flags is that the pretrained weights contain only model weights, and checkpoints, apart from model weights, contain optimizer state, LR scheduler state.

Checkpoints are suitable for dividing the training into parts, for example, in order to divide the training job into shorter stages, or restart training after an infrastructure failure.

Pretrained weights can be used as a base for fine tuning the model to a different dataset, or as a backbone to detection models.

Inference process

Validation is done every epoch, and can be also run separately on a checkpointed model.

python ./main.py --arch efficientnet-<version> --evaluate --epochs 1 --resume <path to checkpoint> -b <batch size> <path to imagenet>

Metrics gathered through training:

• val.loss - validation loss
• val.top1 - validation top1 accuracy
• val.top5 - validation top5 accuracy
• val.total_ips - inference speed measured in images/second
• val.compute_ips - inference speed measured in images/second, not counting data loading
• val.data_time - time spent on waiting on data
• val.compute_time - time spent on inference

To run inference on JPEG image, you have to first extract the model weights from checkpoint:

python checkpoint2model.py --checkpoint-path <path to checkpoint> --weight-path <path where weights will be stored>

Then, run the classification script:

python classify.py --arch efficientnet-<version> --pretrained-from-file <path to weights from previous step> --precision AMP|FP32 --image <path to JPEG image>

You can also run the ImageNet validation on pretrained weights:

python ./main.py --arch efficientnet-<version> --evaluate --epochs 1 --pretrained-from-file <path to pretrained weights> -b <batch size> <path to imagenet>

NGC pretrained weights

Pretrained weights can be downloaded from NGC:

wget <ngc weights url>

URL for each model can be found in the following table:

To run inference on ImageNet, run:

python ./main.py --arch efficientnet-<version> --evaluate --epochs 1 --pretrained -b <batch size> <path to imagenet>

To run inference on JPEG images using pretrained weights, run:

python classify.py --arch efficientnet-<version> --pretrained --precision AMP|FP32 --image <path to JPEG image>

Quantization process

EfficientNet-b0 and EfficientNet-b4 models can be quantized using the QAT process from running the quant_main.py script.

python ./quant_main.py <path to imagenet> --arch efficientnet-quant-<version> --epochs <# of QAT epochs> --pretrained-from-file <path to non-quantized model weights> <any other parameters for training such as batch, momentum etc.>

During the QAT process, evaluation is done in the same way as during standard training. quant_main.py works in the same way as the original main.py script, but with quantized models. It means that quant_main.py can be used to resume the QAT process with the flag --resume:

python ./quant_main.py <path to imagenet> --arch efficientnet-quant-<version> --resume <path to mid-training checkpoint> ...

or to evaluate a created checkpoint with the flag --evaluate:

python ./quant_main.py --arch efficientnet-quant-<version> --evaluate --epochs 1 --resume <path to checkpoint> -b <batch size> <path to imagenet>

It also can run on multi-GPU in an identical way as the standard main.py script:

python ./multiproc.py --nproc_per_node 8 ./quant_main.py --arch efficientnet-quant-<version> ... <path to imagenet>

There is also a possibility to transform trained models (quantized or not) into ONNX format, which is needed to convert it later into TensorRT, where quantized networks are much faster during inference. Conversion to TensorRT will be supported in the next release. The conversion to ONNX consists of two steps:

• translate checkpoint to pure weights:

python checkpoint2model.py --checkpoint-path <path to quant checkpoint> --weight-path <path where quant weights will be stored>

• translate pure weights to ONNX:

python model2onnx.py --arch efficientnet-quant-<version> --pretrained-from-file <path to model quant weights> -b <batch size> --trt True

Quantized models could also be used to classify new images using the classify.py flag. For example: python classify.py --arch efficientnet-quant-<version> --pretrained-from-file <path to quant weights> --image <path to JPEG image>