Quick Start Guide

To train your model using mixed precision with Tensor Cores, perform the following steps using the default parameters of the SSD v1.1 model on the COCO 2017 dataset.

1. Download and preprocess the dataset.

Extract the COCO 2017 dataset with download_dataset.sh $COCO_DIR. Data will be downloaded to the $COCO_DIR directory (on the host).

2. Build the SSD300 v1.1 PyTorch NGC container.

docker build . -t nvidia_ssd

3. Launch the NGC container to run training/inference.

nvidia-docker run --rm -it --ulimit memlock=-1 --ulimit stack=67108864 -v $COCO_DIR:/coco --ipc=host nvidia_ssd

Note: the default mount point in the container is /coco.

4. Start training.

The ./examples directory provides several sample scripts for various GPU settings and act as wrappers around the main.py script. The example scripts need two arguments:

• A path to root SSD directory.
• A path to COCO 2017 dataset.

Remaining arguments are passed to the main.py script.

The --save flag, saves the model after each epoch. The checkpoints are stored as ./models/epoch_*.pt.

Use python main.py -h to obtain the list of available options in the main.py script. For example, if you want to run 8 GPU training with TensorCore acceleration and save checkpoints after each epoch, run:

bash ./examples/SSD300_FP16_8GPU.sh . /coco --save

For information about how to train using mixed precision, see the Mixed Precision Training paper and Training With Mixed Precision documentation.

For PyTorch, easily adding mixed-precision support is available from NVIDIA's APEX, a PyTorch extension that contains utility libraries, such as AMP, which require minimal network code changes to leverage tensor cores performance.

5. Start validation/evaluation.

The main.py training script automatically runs validation during training. The results from the validation are printed to stdout.

Pycocotools' open-sourced scripts provides a consistent way to evaluate models on the COCO dataset. We are using these scripts during validation to measure models performance in AP metric. Metrics below are evaluated using pycocotools' methodology, in the following format:

 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.250
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.423
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.257
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.076
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.269
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.399
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.237
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.342
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.358
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.118
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.394
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.548

The metric reported in our results is present in the first row.

To evaluate a checkpointed model saved in previous point, run:

python ./main.py --backbone resnet50 --mode evaluation --checkpoint ./models/epoch_*.pt --data /coco

6. Optionally, resume training from a checkpointed model.

python ./main.py --backbone resnet50 --checkpoint ./models/epoch_*.pt --data /coco

Details

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

Command line arguments

All these parameters can be controlled by passing command line arguments to the main.py script. To get a complete list of all command line arguments with descriptions and default values you can run:

python main.py --help

Getting the data

The SSD model was trained on the COCO 2017 dataset. The val2017 validation set was used as a validation dataset. PyTorch can work directly on JPEGs, therefore, preprocessing/augmentation is not needed.

This repository contains the download_dataset.sh download script which will automatically download and preprocess the training, validation and test datasets. By default, data will be downloaded to the /coco directory.

Training process

Training the SSD model is implemented in the main.py script.

By default, training is running for 65 epochs. Because evaluation is relatively time consuming, it is not running every epoch. With default settings, evaluation is executed after epochs: 21, 31, 37, 42, 48, 53, 59, 64. The model is evaluated using pycocotools distributed with the COCO dataset. Which epochs should be evaluated can be reconfigured with argument --evaluation.

To run training with Tensor Cores, use the --fp16 flag when running the main.py script. The flag --save flag enables storing checkpoints after each epoch under ./models/epoch_*.pt.

Data preprocessing

Before we feed data to the model, both during training and inference, we perform:

• Normalization
• Encoding bounding boxes
• Resize to 300x300

Data augmentation

During training we perform the following augmentation techniques:

• Random crop
• Random horizontal flip
• Color jitter

Enabling mixed precision

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 the Volta and Turing 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. Manually adding loss scaling to preserve small gradient values.

Mixed precision is enabled in PyTorch by using the Automatic Mixed Precision (AMP), library from APEX that 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 PyTorch, loss scaling can be easily applied by using scale_loss() method provided by amp. The scaling value to be used can be dynamic or fixed.

For an in-depth walk through on AMP, check out sample usage here. APEX is a PyTorch extension that contains utility libraries, such as AMP, which require minimal network code changes to leverage tensor cores performance.

To enable mixed precision, you can:

• Import AMP from APEX, for example:

  from apex import amp

• Initialize an AMP handle, for example:

  amp_handle = amp.init(enabled=True, verbose=True)

• Wrap your optimizer with the AMP handle, for example:

  optimizer = amp_handle.wrap_optimizer(optimizer)

• Scale loss before backpropagation (assuming loss is stored in a variable called losses)

  • Default backpropagate for FP32:

    losses.backward()

  • Scale loss and backpropagate with AMP:

    with optimizer.scale_loss(losses) as scaled_losses:
       scaled_losses.backward()
    

For information about:

• How to train using mixed precision, see the Mixed Precision Training paper and Training With Mixed Precision documentation.
• Techniques used for mixed precision training, see the Mixed-Precision Training of Deep Neural Networks blog.