Skip to content

Commit a02d310

Browse files
YosuaMichaelfacebook-github-bot
authored andcommitted
[fbsync] Add MaskRCNN improved weights (#5773)
Summary: * Add MaskRCNN improved weights * Adding recipe URL (Note: this ignores all push blocking failures!) Reviewed By: jdsgomes, NicolasHug Differential Revision: D36095710 fbshipit-source-id: ad917bbd12f545f62681bb4607b74b2ebe650592
1 parent 345b597 commit a02d310

File tree

1 file changed

+22
-5
lines changed

1 file changed

+22
-5
lines changed

torchvision/models/detection/mask_rcnn.py

Lines changed: 22 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -349,17 +349,22 @@ def __init__(self, in_channels, dim_reduced, num_classes):
349349
# nn.init.constant_(param, 0)
350350

351351

352+
_COMMON_META = {
353+
"task": "image_object_detection",
354+
"architecture": "MaskRCNN",
355+
"categories": _COCO_CATEGORIES,
356+
"interpolation": InterpolationMode.BILINEAR,
357+
}
358+
359+
352360
class MaskRCNN_ResNet50_FPN_Weights(WeightsEnum):
353361
COCO_V1 = Weights(
354362
url="https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth",
355363
transforms=ObjectDetection,
356364
meta={
357-
"task": "image_object_detection",
358-
"architecture": "MaskRCNN",
365+
**_COMMON_META,
359366
"publication_year": 2017,
360367
"num_params": 44401393,
361-
"categories": _COCO_CATEGORIES,
362-
"interpolation": InterpolationMode.BILINEAR,
363368
"recipe": "https://github.com/pytorch/vision/tree/main/references/detection#mask-r-cnn",
364369
"map": 37.9,
365370
"map_mask": 34.6,
@@ -369,7 +374,19 @@ class MaskRCNN_ResNet50_FPN_Weights(WeightsEnum):
369374

370375

371376
class MaskRCNN_ResNet50_FPN_V2_Weights(WeightsEnum):
372-
pass
377+
COCO_V1 = Weights(
378+
url="https://download.pytorch.org/models/maskrcnn_resnet50_fpn_v2_coco-73cbd019.pth",
379+
transforms=ObjectDetection,
380+
meta={
381+
**_COMMON_META,
382+
"publication_year": 2021,
383+
"num_params": 46359409,
384+
"recipe": "https://github.com/pytorch/vision/pull/5773",
385+
"map": 47.4,
386+
"map_mask": 41.8,
387+
},
388+
)
389+
DEFAULT = COCO_V1
373390

374391

375392
@handle_legacy_interface(

0 commit comments

Comments
 (0)