Add RetinaNet improved weights by datumbox · Pull Request #5756 · pytorch/vision (original) (raw)

Fixes #5307

Adds new pre-trained weights for RetinaNet + ResNet50 + FPN for the v2 variant with post-paper optimizations (no FrozenBN + c5 instead of p5 input on extra layers + GN on Head + gIoU loss). It improves the previous baseline by +5.1 mAP.

Trained with:

python -u run_with_submitit.py --ngpus 8 --nodes 1 --weights-backbone ResNet50_Weights.IMAGENET1K_V2 \ 
--dataset coco --model retinanet_resnet50_fpn_v2 --opt adamw --lr 0.0001 --epochs 26 --lr-steps 16 22 \
--weight-decay 0.05 --norm-weight-decay 0.0 --data-augmentation multiscale --sync-bn

Verified with:

torchrun --nproc_per_node=1 train.py --test-only --weights RetinaNet_ResNet50_FPN_V2_Weights.COCO_V1 \
--model retinanet_resnet50_fpn_v2 -b 1

IoU metric: bbox
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.415
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.618
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.439
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.270
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.457
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.538
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.337
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.543
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.587
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.418
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.629
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.721