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

Scripts and sample code

Descriptions of the key scripts and folders are provided below.

• model and object_detection - Contains code to build individual components of the model such as backbone, FPN, RPN, classification and bbox heads, and so on.

• data - Contains code to convert raw data to tfrecords

• dataset - Contains code to build the data pipeline such as dataloader, transforms, dataset builder.

• scripts/ - Contains shell scripts to launch training and evaluation of the model and perform inference.

  • D0/convergence-{AMP, TF32, FP32}-{8, 32}x{V100-32G, A100-80G}.sh - Launches model training

  • D0/evaluate-{AMP, FP32, TF32}-{8, 32}x{A100-80G, V100-32G}.sh - Performs inference and computes mAP of predictions.

  • D0/inference.sh - Performs inference on an image

  • D0/training-benchmark-{AMP, TF32, FP32}-{V100-32G, A100-80G}.sh - Launches model training for 500 iterations to benchmark training

  • D0/inference-benchmark.sh - Benchmarks inference

  • docker/ - Scripts to build the docker image and to start an interactive session

• utils/

  • Contains utility components like default hyper parameters, checkpoint utils, training utils, and so on.

• train.py - End to end to script to load data, build and train the model.

• eval.py - End to end script to load data and checkpoint to perform inference and compute mAP scores.

Parameters

train.py script parameters

Important parameters for training are listed below with defaults.

Command-line options

To display the full list of available options and their descriptions, use the --helpshort command-line option:

  --amp:                Enable mixed precision training
  --backbone_init:      Initialize backbone weights from 
                        checkpoint in this directory
  --batch_size:         training local batch size
  --benchmark:          Train for a fixed number of steps for performance
  --benchmark_steps:    Train for these many steps to benchmark training
    performance
  --checkpoint_period:  Save ema model weights after every X epochs for eval
  --debug:              Enable debug mode
  --enable_map_parallelization: 
                        Parallelize stateless map transformations in dataloader
  --eval_samples:       The number of samples for evaluation.
  --hparams:            Comma separated k=v pairs of hyperparameters or a module 
                        containing attributes to use as hyperparameters
  --log_filename:       Filename for dllogger logs
  --log_steps:          Interval of steps between logging of batch level stats
  --lr:                 Learning rate
  --lr_tb:              Log learning rate at each step to TB
  --training_mode:      Training mode (train/traineval/train300)
  --model_dir:          Location of model_dir
  --model_name:         Model name
  --num_epochs:         Number of epochs for training
  --num_examples_per_epoch: 
                        Number of examples in one epoch (coco is 117266).
                        Default is 120000.
  --pretrained_ckpt: Start training from this EfficientDet checkpoint.
  --seed:               Random seed
  --set_num_threads:    Set inter-op and intra-op parallelism threads
  --testdev_dir:        COCO testdev dir. If not None, ignore val_json_file.
  --time_history:       Get time history
  --training_file_pattern: 
                        Glob for training data files, e.g., coco/train-*.tfrecord.
  --use_fake_data:      Use fake input
  --use_xla:            Use XLA
  --val_file_pattern:   Glob for eval tfrecords, e.g. coco/val-*.tfrecord.
  --val_json_file:      COCO validation JSON containing golden bounding boxes. If
                        None, use the ground truth from the dataloader. Ignored if 
                        testdev_dir is not None.
  --validate:           Get validation loss after each epoch
  --warmup_epochs:      Number of warmup epochs
  --warmup_value:       Initial warmup value

The default training_mode (traineval) is training along with evaluation. Note that evaluation starts only after 200 epochs of training, and the frequency of evaluation can be set by setting checkpoint_period=<every n epochs> which is currently set to 10. Also, in the traineval training_mode, the model stops training at 300 epochs to avoid overfitting. To run only training without evaluation, set the training_mode to train. In this training_mode, the ema checkpoints are stored in path model_dir/ema_weights/ every checkpoint_period (in epochs) . The training stops after training for 75% of the total number of epochs, and the last ema-weight checkpoint is evaluated. For benchmarking only the training time for 300 epochs, the training_mode can be set to train300 where the model trains for exactly 300 epochs without any evaluation.

Getting the data

By default, the EfficientDet model is trained on the COCO 2017 dataset. This dataset comes with a training and validation set. Follow steps in the Quick Start Guide to download and pre-process the dataset into tfrecord format.

Training Process

Training is performed using the train.py script. The default parameters can be overridden by command-line arguments.

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

By default, bash script scripts/D0/convergence-{AMP, FP32, TF32}-8x{A100-80G, V100-32G}.sh will start the training process from scratch with the following settings.

• Use 8 GPUs
• Saves checkpoints after every 10 epochs to model_dir which is /tmp/convergence-{AMP, FP32, TF32}-8x{A100-80G, V100-32G} folder

To resume from a checkpoint, just make sure that the model_dir stays the same and that the checkpoints saved are already present in it.

Multi-node

Multi-node runs can be launched on a Pyxis/enroot Slurm cluster (refer to Requirements) with the ./scripts/D0/convergence-{AMP, FP32}-32xV100-32G.sub script with the following command for a 4-node NVIDIA DGX V100 example:

checkpointdir=<path to efficientnet B0 pretrained checkpoint directory> datadir=<path to coco 2017 dataset in tfrecord format> sbatch N 4 --ntasks-per-node=8 --ntasks=32 ./scripts/D0/convergence-{AMP, FP32}-32xV100-32G.sub

Note that the ./scripts/D0/convergence-{AMP, FP32}-32xV100-32G.sub script is a starting point that has to be adapted depending on the environment. In particular, variables such as --container-image handles the container image to train using, datadir handles the location of the COCO-2017 data, and checkpointdir has the path to the pre-trained backbone (EfficientNet) weights.

Refer to the file's content to view the full list of variables to adjust for your system.