Relation Networks for Object Detection based on Pytorch and Torchvision

This is a Faster-RCNN with FPN and relation network based on Torchvision.

This folder contains the training scripts for object detection, Faster-RCNN with FPN model, Relation Network and evaluation scripts.

Progress

add Relation Network between two fully connected layers in RoI head
end to end train this model on COCO 2014
evaluate mAP based on COCO evaluation tools
replace NMS by Relation Network
improve the performance (Current mAP: 36)

Requirements

Python 3.6
Pytorch 1.4.1 and Torchvision
The Python packages:

cython
pycocotools
matplotlib

Dataset

Please download the COCO Dataset.

Before Training

You must modify the following flags:

--data-path=/path/to/coco/dataset

--nproc_per_node=<number_of_gpus_available>

Data-path is the path of dataset.

nproc_per_node is the number of GPUs.

Training

Train script is below, the backbone is Resnet 101:

python -m torch.distributed.launch --nproc_per_node=4 --use_env train.py\
    --dataset coco --model fasterrcnn_resnet101_fpn --epochs 30\
    --lr-steps 16 22 --aspect-ratio-group-factor 3

This model contains {1, 1} relation network between the fully connected layers in RoI head.