Skip to content

Repository for Single Shot MultiBox Detector and its variants, implemented with pytorch, python3.

License

Notifications You must be signed in to change notification settings

ShuangXieIrene/ssds.pytorch

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

69 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

ssds.pytorch

Repository for Single Shot MultiBox Detector and its variants, implemented with pytorch, python3. This repo is easy to setup and has plenty of visualization methods. We hope this repo can help people have a better understanding for ssd-like model and help people train and deploy the ssds model easily.

Currently, it contains these features:

  • Multiple SSD Variants: ssd, fpn, bifpn, yolo and etc.
  • Multiple Base Network: resnet, regnet, mobilenet and etc.
  • Visualize the features of the ssd-like models to help the user understand the model design and performance.
  • Fast Training and Inference: Utilize Nvidia Apex and Dali to fast training and support the user convert the model to ONNX or TensorRT for deployment.

This repo is depended on the work of ODTK, Detectron and Tensorflow Object Detection API. Thanks for their works.

Notice The pretrain model for the current version does not finished yet, please check the previous version for enrich pretrain models.

Table of Contents

Installation

requirements

  • python>=3.7
  • CUDA>=10.0
  • pytorch>=1.4

basic installation:

conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
git clone https://github.com/ShuangXieIrene/ssds.pytorch.git
cd ssds.pytorch
python setup.py clean -a install

extra python libs for parallel training

Currently, nvidia DALI and apex is not include in the requirements.txt and need to install manually.

pip install --extra-index-url https://developer.download.nvidia.com/compute/redist/cuda/10.0 nvidia-dali
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

Docker

git clone https://github.com/ShuangXieIrene/ssds.pytorch.git
docker build -t ssds:local ./ssds.pytorch/
docker run --gpus all -it --rm -v /data:/data ssds:local

Usage

0. Check the config file by Visualization

Defined the network in a config file and tweak the config file based on the visualized anchor boxes

python -m ssds.utils.visualize -cfg experiments/cfgs/tests/test.yml

1. Training

# basic training
python -m ssds.utils.train -cfg experiments/cfgs/tests/test.yml
# parallel training
python -m torch.distributed.launch --nproc_per_node={num_gpus} -m ssds.utils.train_ddp -cfg experiments/cfgs/tests/test.yml

2. Evaluation

python -m ssds.utils.train -cfg experiments/cfgs/tests/test.yml -e

3. Export to ONNX or TRT model

python -m ssds.utils.export -cfg experiments/cfgs/tests/test.yml -c best_mAP.pth -h

Performance

Visualization

About

Repository for Single Shot MultiBox Detector and its variants, implemented with pytorch, python3.

Topics

Resources

License

Stars

Watchers

Forks

Packages

No packages published