With a ResNet-50 backbone and a number of architectural modifications, this version provides better accuracy and performance.
To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the SSD v1.1 model on the COCO 2017 dataset. For the specifics concerning training and inference, see the Advanced section.
- Clone the repository.
git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/Detection/SSD
- 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).
- Build the SSD300 v1.1 PyTorch NGC container.
docker build . -t nvidia_ssd
- Start an interactive session in the NGC container to run training/inference.
docker run --rm -it --gpus=all --ipc=host -v $COCO_DIR:/coco nvidia_ssd
Note: the default mount point in the container is /coco.
- 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 the root SSD directory.
- A path to the COCO 2017 dataset.
Remaining arguments are passed to the main.py script.
The --save save_dir flag, saves the model after each epoch in save_dir directory.
The checkpoints are stored as <save_dir>/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 Tensor Core acceleration and
save checkpoints after each epoch, run:
bash ./examples/SSD300_FP16_8GPU.sh . /coco --save $SSD_CHECKPINT_PATH
- Start validation/evaluation.
The main.py training script automatically runs validation during training.
The results from the validation are printed to stdout.
To evaluate a checkpointed model saved in the previous point, run:
python ./main.py --backbone resnet50 --mode evaluation --checkpoint ./models/epoch_*.pt --data /coco
- Optionally, resume training from a checkpointed model.
python ./main.py --backbone resnet50 --checkpoint ./models/epoch_*.pt --data /coco
- Start inference/predictions.
You can check your trained model with a Jupyter notebook provided in the examples directory. Start with running a Docker container with a Jupyter notebook server:
docker run --rm -it --gpus=all --ipc=host -v $SSD_CHECKPOINT_PATH:/checkpoints/SSD300v1.1.pt -v $COCO_PATH:/datasets/coco2017 -p 8888:8888 nvidia_ssd jupyter-notebook --ip 0.0.0.0 --allow-root