[fbsync] Improved meta-data for models (#5170)

Summary:
* Improved meta-data for models.

* Addressing comments from code-review.

* Add parameter count.

* Fix linter.

Reviewed By: sallysyw

Differential Revision: D33479281

fbshipit-source-id: 7a133324ed5a289a0ac89522b0d4a38ce8b201e0

Author: prabhat00155

test/test_prototype_models.py

@@ -94,29 +94,47 @@ def test_naming_conventions(model_fn):
     + TM.get_models_from_module(models.video)
     + TM.get_models_from_module(models.optical_flow),
 )
+@run_if_test_with_prototype
 def test_schema_meta_validation(model_fn):
     classification_fields = ["size", "categories", "acc@1", "acc@5"]
     defaults = {
-        "all": ["interpolation", "recipe"],
+        "all": ["task", "architecture", "publication_year", "interpolation", "recipe", "num_params"],
         "models": classification_fields,
         "detection": ["categories", "map"],
         "quantization": classification_fields + ["backend", "quantization", "unquantized"],
         "segmentation": ["categories", "mIoU", "acc"],
         "video": classification_fields,
         "optical_flow": [],
     }
+    model_name = model_fn.__name__
     module_name = model_fn.__module__.split(".")[-2]
     fields = set(defaults["all"] + defaults[module_name])
 
     weights_enum = _get_model_weights(model_fn)
+    if len(weights_enum) == 0:
+        pytest.skip(f"Model '{model_name}' doesn't have any pre-trained weights.")
 
     problematic_weights = {}
+    incorrect_params = []
     for w in weights_enum:
         missing_fields = fields - set(w.meta.keys())
         if missing_fields:
             problematic_weights[w] = missing_fields
+        if w == weights_enum.default:
+            if module_name == "quantization":
+                # parametes() cound doesn't work well with quantization, so we check against the non-quantized
+                unquantized_w = w.meta.get("unquantized")
+                if unquantized_w is not None and w.meta.get("num_params") != unquantized_w.meta.get("num_params"):
+                    incorrect_params.append(w)
+            else:
+                if w.meta.get("num_params") != sum(p.numel() for p in model_fn(weights=w).parameters()):
+                    incorrect_params.append(w)
+        else:
+            if w.meta.get("num_params") != weights_enum.default.meta.get("num_params"):
+                incorrect_params.append(w)
 
     assert not problematic_weights
+    assert not incorrect_params
 
 
 @pytest.mark.parametrize("model_fn", TM.get_models_from_module(models))

torchvision/models/segmentation/fcn.py

@@ -18,7 +18,9 @@
 
 class FCN(_SimpleSegmentationModel):
     """
-    Implements a Fully-Convolutional Network for semantic segmentation.
+    Implements FCN model from
+    `"Fully Convolutional Networks for Semantic Segmentation"
+    <https://arxiv.org/abs/1411.4038>`_.
 
     Args:
         backbone (nn.Module): the network used to compute the features for the model.

torchvision/prototype/models/alexnet.py

@@ -18,6 +18,10 @@ class AlexNet_Weights(WeightsEnum):
         url="https://download.pytorch.org/models/alexnet-owt-7be5be79.pth",
         transforms=partial(ImageNetEval, crop_size=224),
         meta={
+            "task": "image_classification",
+            "architecture": "AlexNet",
+            "publication_year": 2012,
+            "num_params": 61100840,
             "size": (224, 224),
             "categories": _IMAGENET_CATEGORIES,
             "interpolation": InterpolationMode.BILINEAR,

torchvision/prototype/models/densenet.py

@@ -64,6 +64,9 @@ def _densenet(
 
 
 _COMMON_META = {
+    "task": "image_classification",
+    "architecture": "DenseNet",
+    "publication_year": 2016,
     "size": (224, 224),
     "categories": _IMAGENET_CATEGORIES,
     "interpolation": InterpolationMode.BILINEAR,
@@ -77,6 +80,7 @@ class DenseNet121_Weights(WeightsEnum):
         transforms=partial(ImageNetEval, crop_size=224),
         meta={
             **_COMMON_META,
+            "num_params": 7978856,
             "acc@1": 74.434,
             "acc@5": 91.972,
         },
@@ -90,6 +94,7 @@ class DenseNet161_Weights(WeightsEnum):
         transforms=partial(ImageNetEval, crop_size=224),
         meta={
             **_COMMON_META,
+            "num_params": 28681000,
             "acc@1": 77.138,
             "acc@5": 93.560,
         },
@@ -103,6 +108,7 @@ class DenseNet169_Weights(WeightsEnum):
         transforms=partial(ImageNetEval, crop_size=224),
         meta={
             **_COMMON_META,
+            "num_params": 14149480,
             "acc@1": 75.600,
             "acc@5": 92.806,
         },
@@ -116,6 +122,7 @@ class DenseNet201_Weights(WeightsEnum):
         transforms=partial(ImageNetEval, crop_size=224),
         meta={
             **_COMMON_META,
+            "num_params": 20013928,
             "acc@1": 76.896,
             "acc@5": 93.370,
         },

torchvision/prototype/models/detection/faster_rcnn.py

@@ -31,6 +31,9 @@
 
 
 _COMMON_META = {
+    "task": "image_object_detection",
+    "architecture": "FasterRCNN",
+    "publication_year": 2015,
     "categories": _COCO_CATEGORIES,
     "interpolation": InterpolationMode.BILINEAR,
 }
@@ -42,6 +45,7 @@ class FasterRCNN_ResNet50_FPN_Weights(WeightsEnum):
         transforms=CocoEval,
         meta={
             **_COMMON_META,
+            "num_params": 41755286,
             "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-resnet-50-fpn",
             "map": 37.0,
         },
@@ -55,6 +59,7 @@ class FasterRCNN_MobileNet_V3_Large_FPN_Weights(WeightsEnum):
         transforms=CocoEval,
         meta={
             **_COMMON_META,
+            "num_params": 19386354,
             "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-mobilenetv3-large-fpn",
             "map": 32.8,
         },
@@ -68,6 +73,7 @@ class FasterRCNN_MobileNet_V3_Large_320_FPN_Weights(WeightsEnum):
         transforms=CocoEval,
         meta={
             **_COMMON_META,
+            "num_params": 19386354,
             "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#faster-r-cnn-mobilenetv3-large-320-fpn",
             "map": 22.8,
         },

torchvision/prototype/models/detection/keypoint_rcnn.py

@@ -24,6 +24,9 @@
 
 
 _COMMON_META = {
+    "task": "image_object_detection",
+    "architecture": "KeypointRCNN",
+    "publication_year": 2017,
     "categories": _COCO_PERSON_CATEGORIES,
     "keypoint_names": _COCO_PERSON_KEYPOINT_NAMES,
     "interpolation": InterpolationMode.BILINEAR,
@@ -36,6 +39,7 @@ class KeypointRCNN_ResNet50_FPN_Weights(WeightsEnum):
         transforms=CocoEval,
         meta={
             **_COMMON_META,
+            "num_params": 59137258,
             "recipe": "https://github.com/pytorch/vision/issues/1606",
             "map": 50.6,
             "map_kp": 61.1,
@@ -46,6 +50,7 @@ class KeypointRCNN_ResNet50_FPN_Weights(WeightsEnum):
         transforms=CocoEval,
         meta={
             **_COMMON_META,
+            "num_params": 59137258,
             "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#keypoint-r-cnn",
             "map": 54.6,
             "map_kp": 65.0,

torchvision/prototype/models/detection/mask_rcnn.py

@@ -28,6 +28,10 @@ class MaskRCNN_ResNet50_FPN_Weights(WeightsEnum):
         url="https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth",
         transforms=CocoEval,
         meta={
+            "task": "image_object_detection",
+            "architecture": "MaskRCNN",
+            "publication_year": 2017,
+            "num_params": 44401393,
             "categories": _COCO_CATEGORIES,
             "interpolation": InterpolationMode.BILINEAR,
             "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#mask-r-cnn",

torchvision/prototype/models/detection/retinanet.py

@@ -29,6 +29,10 @@ class RetinaNet_ResNet50_FPN_Weights(WeightsEnum):
         url="https://download.pytorch.org/models/retinanet_resnet50_fpn_coco-eeacb38b.pth",
         transforms=CocoEval,
         meta={
+            "task": "image_object_detection",
+            "architecture": "RetinaNet",
+            "publication_year": 2017,
+            "num_params": 34014999,
             "categories": _COCO_CATEGORIES,
             "interpolation": InterpolationMode.BILINEAR,
             "recipe": "https://github.com/pytorch/vision/tree/main/references/detection#retinanet",

torchvision/prototype/models/detection/ssd.py

@@ -27,6 +27,10 @@ class SSD300_VGG16_Weights(WeightsEnum):
         url="https://download.pytorch.org/models/ssd300_vgg16_coco-b556d3b4.pth",
         transforms=CocoEval,
         meta={
+            "task": "image_object_detection",
+            "architecture": "SSD",
+            "publication_year": 2015,
+            "num_params": 35641826,
             "size": (300, 300),
             "categories": _COCO_CATEGORIES,
             "interpolation": InterpolationMode.BILINEAR,

torchvision/prototype/models/detection/ssdlite.py

@@ -32,6 +32,10 @@ class SSDLite320_MobileNet_V3_Large_Weights(WeightsEnum):
         url="https://download.pytorch.org/models/ssdlite320_mobilenet_v3_large_coco-a79551df.pth",
         transforms=CocoEval,
         meta={
+            "task": "image_object_detection",
+            "architecture": "SSDLite",
+            "publication_year": 2018,
+            "num_params": 3440060,
             "size": (320, 320),
             "categories": _COCO_CATEGORIES,
             "interpolation": InterpolationMode.BILINEAR,