SSD in a single pytorch file

The data loader and augmentation are directly copied from Max deGroot.

The SSD model is implemented in one file: ssd.py.

My result on VOC Dataset: mAP=0.7635

Datasets

To make things easy, we provide bash scripts to handle the dataset downloads and setup for you. We also provide simple dataset loaders that inherit torch.utils.data.Dataset, making them fully compatible with the torchvision.datasets API.

COCO

Microsoft COCO: Common Objects in Context

Download COCO 2014

# specify a directory for dataset to be downloaded into, else default is ~/data/
sh data/scripts/COCO2014.sh

VOC Dataset

PASCAL VOC: Visual Object Classes

Download VOC2007 trainval & test

# specify a directory for dataset to be downloaded into, else default is ~/data/
sh data/scripts/VOC2007.sh # <directory>

Training SSD

First download the fc-reduced VGG-16 PyTorch base network weights at: https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth
By default, we assume you have downloaded the file in the ssd.pytorch/weights dir:

mkdir weights
cd weights
wget https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth

To train SSD using the train script simply specify the parameters listed in train.py as a flag or manually change them.

python train.py

Note:
- For training, an NVIDIA GPU is strongly recommended for speed.
- You can pick-up training from a checkpoint by specifying the path as one of the training parameters (again, see train.py for options)

Evaluation

To evaluate a trained network:

python eval.py

References

Wei Liu, et al. "SSD: Single Shot MultiBox Detector." ECCV2016.
Original Implementation (CAFFE)
Pytorch version by Alex Koltun.

iamhankai / ssd-pytorch