- Python 3.X
- Pytorch >1.0.0
- OpenCV
- numpy
- matplotlib
- CUDA (tested on CUDA 9, CUDA 10)
- cuDNN
python create_vgg16_weights.py
python train.py --help
usage: train.py [-h] [--data_dir DATA_DIR] [--augmentation AUGMENTATION]
[--pretrained_path PRETRAINED_PATH] [--neg_ratio NEG_RATIO]
[--lr LR] [--momentum MOMENTUM] [--weight_decay WEIGHT_DECAY]
[--gamma GAMMA] [--arch ARCH] [--num_examples NUM_EXAMPLES]
[--batch_size BATCH_SIZE] [--num_epochs NUM_EPOCHS]
[--checkpoint_dir CHECKPOINT_DIR]
optional arguments:
-h, --help show this help message and exit
--data_dir DATA_DIR data root directory
--augmentation AUGMENTATION
whether to perform augmentation of data
--pretrained_path PRETRAINED_PATH
pretrained weight path
--neg_ratio NEG_RATIO
negative examples ratio (used in hard negative mining
logic)
--lr LR initial learning rate
--momentum MOMENTUM momentum (used in SGD)
--weight_decay WEIGHT_DECAY
weight decay (used in SGD)
--gamma GAMMA gamma (used in SGD)
--arch ARCH SSD architecture, can be either ssd300 or ssd512
--num_examples NUM_EXAMPLES
number of examples to train
--batch_size BATCH_SIZE
batch size
--num_epochs NUM_EPOCHS
number of epochs to train
--checkpoint_dir CHECKPOINT_DIR
directory to store training weights
Example 1: train SSD300 on 40 examples, batch_size=10 without augmentation using the pretrained weights generated above
python train.py --num_examples 40 --batch_size 10 --pretrained_path ./weights/new_vgg.pth --augmentation False
Example 2: train SSD512 on full dataset, batch_size=32 using the pretrained weights generated above on 120 epochs
python train.py --pretrained_path ./weights/new_vgg.pth --arch ssd512 --num_epochs 120
python test.py --help
usage: test.py [-h] [--data_dir DATA_DIR] [--arch ARCH]
[--save_image_dir SAVE_IMAGE_DIR]
[--pretrained_path PRETRAINED_PATH]
[--num_examples NUM_EXAMPLES]
[--max_num_boxes_per_class MAX_NUM_BOXES_PER_CLASS]
[--score_thresh SCORE_THRESH] [--nms_thresh NMS_THRESH]
[--batch_size BATCH_SIZE]
optional arguments:
-h, --help show this help message and exit
--data_dir DATA_DIR data root directory
--arch ARCH SSD architecture, can be either ssd300 or ssd512
--save_image_dir SAVE_IMAGE_DIR
directory to save result images
--pretrained_path PRETRAINED_PATH
pretrained weight path
--num_examples NUM_EXAMPLES
number of examples to test
--max_num_boxes_per_class MAX_NUM_BOXES_PER_CLASS
maximum number of objects per class (used in NMS
logic)
--score_thresh SCORE_THRESH
score threshold
--nms_thresh NMS_THRESH
NMS threshold
--batch_size BATCH_SIZE
batch size
Example: test the SSD512 trained by the command above on full dataset
python test.py --pretrained_path ./models/ssd_epoch_120.pth --arch ssd512 --num_examples -1
- Random patching
- Random distortion
- VOC Evaluation
- OpenCV -> PIL for better font setting
SSD: Single Shot MultiBox Detector