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

Scripts and sample code

In the root directory, the most important files are:

• main.py: the script that controls the logic of training and validation of the SSD300 v1.1 model;
• Dockerfile: Instructions for docker to build a container with the basic set of dependencies to run SSD300 v1.1;
• requirements.txt: a set of extra Python requirements for running SSD300 v1.1;
• download_dataset.py: automatically downloads the COCO dataset for training.

The ssd/ directory contains modules used to train and evaluate the SSD300 v1.1 model

• model.py: the definition of SSD300 v1.1 model
• data.py: definition of input pipelines used in training and evaluation
• train.py: functions used to train the SSD300 v1.1 model
• evaluate.py: functions used to evaluate the SSD300 v1.1 model
• coco_pipeline.py: definition of input pipeline using NVIDIA DALI
• coco.py: code specific for the COCO dataset
• logger.py: utilities for logging
• utils.py: extra utility functions

The examples/ directory contains scripts wrapping common scenarios.

Parameters

The script main.py

The script for training end evaluating the SSD300 v1.1 model have a variety of parameters that control these processes.

Common parameters

--data : use it to specify, where your dataset is. By default, the script will look for it under the /coco directory.

--checkpoint : allows you to specify the path to the pre-trained model.

--save : when the flag is turned on, the script will save the trained model checkpoints in the specified directory

--seed : Use it to specify the seed for RNGs.

--amp : when the flag is turned on, the AMP features will be enabled.

Training related

--epochs : a number of times the model will see every example from the training dataset.

--evaluation : after this parameter, list the number of epochs after which evaluation should be performed.

--learning-rate : initial learning rate.

--multistep : after this parameter, list the epochs after which learning rate should be decayed.

--warmup : allows you to specify the number of iterations for which a linear learning-rate warmup will be performed.

--momentum : momentum argument for SGD optimizer.

--weight-decay : weight decay argument for SGD optimizer.

--batch-size : a number of inputs processed at once for each iteration.

--backbone-path : the path to the checkpointed backbone. When it is not provided, a pre-trained model from torchvision will be downloaded.

Evaluation related

--eval-batch-size : a number of inputs processed at once for each iteration.

Utility parameters

--help : displays a short description of all parameters accepted by the script.

Command-line options

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.

Dataset guidelines

Our model expects input data aligned in a way a COCO dataset is aligned by the download_dataset.sh script. train2017 and val2017 directories should contain images in JPEG format. Annotation format is described in the COCO documentation.

The preprocessing of the data is defined in the ssd/coco_pipeline.py module.

Data preprocessing

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

• JPEG decoding
• normalization with a mean = [0.485, 0.456, 0.406] and std dev = [0.229, 0.224, 0.225]
• encoding bounding boxes
• resizing to 300x300

Additionally, during training, data is:

• randomly shuffled
• samples without annotations are skipped

Data augmentation

During training we perform the following augmentation techniques:

• Random crop using the algorithm described in the SSD: Single Shot MultiBox Detector paper
• Random horizontal flip
• Color jitter

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 the --evaluation argument.

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

Evaluation process

Pycocotools' open-sourced scripts provides a consistent way to evaluate models on the COCO dataset. We are using these scripts during validation to measure a 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.27205
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.45869
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.27884
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.08275
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.29840
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.42722
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.25092
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.36528
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.38262
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.13577
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.42287
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.57277

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

Inference process

Our scripts for SSD300 v1.1 presents two ways to run inference. To get meaningful results, you need a pre-trained model checkpoint.

One way is to run an interactive session on Jupyter notebook, as described in a 8th step of the Quick Start Guide.

The container prints Jupyter notebook logs like this:

[I 16:17:58.935 NotebookApp] Writing notebook server cookie secret to /root/.local/share/jupyter/runtime/notebook_cookie_secret
[I 16:17:59.769 NotebookApp] JupyterLab extension loaded from /opt/conda/lib/python3.6/site-packages/jupyterlab
[I 16:17:59.769 NotebookApp] JupyterLab application directory is /opt/conda/share/jupyter/lab
[I 16:17:59.770 NotebookApp] Serving notebooks from local directory: /workspace
[I 16:17:59.770 NotebookApp] The Jupyter Notebook is running at:
[I 16:17:59.770 NotebookApp] http://(65935d756c71 or 127.0.0.1):8888/?token=04c78049c67f45a4d759c8f6ddd0b2c28ac4eab60d81be4e
[I 16:17:59.770 NotebookApp] Use Control-C to stop this server and shut down all kernels (twice to skip confirmation).
[W 16:17:59.774 NotebookApp] No web browser found: could not locate runnable browser.
[C 16:17:59.774 NotebookApp]

    To access the notebook, open this file in a browser:
        file:///root/.local/share/jupyter/runtime/nbserver-1-open.html
    Or copy and paste one of these URLs:
        http://(65935d756c71 or 127.0.0.1):8888/?token=04c78049c67f45a4d759c8f6ddd0b2c28ac4eab60d81be4e

Use the token printed in the last line to start your notebook session. The notebook is in examples/inference.ipynb, for example:

http://127.0.0.1:8888/notebooks/examples/inference.ipynb?token=04c78049c67f45a4d759c8f6ddd0b2c28ac4eab60d81be4e

Another way is to run a script examples/SSD300_inference.py. It contains the logic from the notebook, wrapped into a Python script. The script contains sample usage.

To use the inference example script in your own code, you can call the main function, providing input image URIs as an argument. The result will be a list of detections for each input image.