Adding FLOPs and size to model metadata

docs/source/conf.py

@@ -385,27 +385,41 @@ def generate_weights_table(module, table_name, metrics, dataset, include_pattern
     if exclude_patterns is not None:
         weights = [w for w in weights if all(p not in str(w) for p in exclude_patterns)]
 
+    rich_metadata = ["Size (MB)"]
+    if "_flops" in weights[0].meta:  # assumes same rich meta for all models in module
+        rich_metadata = ["GFLOPs"] + rich_metadata
+
     metrics_keys, metrics_names = zip(*metrics)
-    column_names = ["Weight"] + list(metrics_names) + ["Params", "Recipe"]
+    column_names = ["Weight"] + list(metrics_names) + ["Params"] + rich_metadata + ["Recipe"]  # Final column order
     column_names = [f"**{name}**" for name in column_names]  # Add bold
 
-    content = [
-        (
+    content = []
+    for w in weights:
+        row = [
             f":class:`{w} <{type(w).__name__}>`",
             *(w.meta["_metrics"][dataset][metric] for metric in metrics_keys),
             f"{w.meta['num_params']/1e6:.1f}M",
-            f"`link <{w.meta['recipe']}>`__",
-        )
-        for w in weights
-    ]
+        ]
+
+        if "_flops" in w.meta:
+            row.append(f"{w.meta['_flops']:.3f}")
+
+        row.append(f"{round(w.meta['_weight_size'], 1):.1f}")
+        row.append(f"`link <{w.meta['recipe']}>`__")
+
+        content.append(row)
+
+    column_widths = ["110"] + ["18"] * len(metrics_names) + ["18"] + ["18"] * len(rich_metadata) + ["10"]
+    widths_table = " ".join(column_widths)
+
     table = tabulate(content, headers=column_names, tablefmt="rst")
 
     generated_dir = Path("generated")
     generated_dir.mkdir(exist_ok=True)
     with open(generated_dir / f"{table_name}_table.rst", "w+") as table_file:
         table_file.write(".. rst-class:: table-weights\n")  # Custom CSS class, see custom_torchvision.css
         table_file.write(".. table::\n")
-        table_file.write(f"    :widths: 100 {'20 ' * len(metrics_names)} 20 10\n\n")
+        table_file.write(f"    :widths: {widths_table} \n\n")
         table_file.write(f"{textwrap.indent(table, ' ' * 4)}\n\n")
 
 

test/test_extended_models.py

@@ -155,21 +155,26 @@ def test_schema_meta_validation(model_fn):
         "recipe",
         "unquantized",
         "_docs",
+        "_flops",
+        "_weight_size",
     }
     # mandatory fields for each computer vision task
     classification_fields = {"categories", ("_metrics", "ImageNet-1K", "acc@1"), ("_metrics", "ImageNet-1K", "acc@5")}
+    rich_metadata = {"_flops"}
     defaults = {
-        "all": {"_metrics", "min_size", "num_params", "recipe", "_docs"},
-        "models": classification_fields,
-        "detection": {"categories", ("_metrics", "COCO-val2017", "box_map")},
+        "all": {"_metrics", "min_size", "num_params", "recipe", "_docs", "_weight_size"},
+        "models": classification_fields | rich_metadata,
+        "detection": {"categories", ("_metrics", "COCO-val2017", "box_map")} | rich_metadata,
         "quantization": classification_fields | {"backend", "unquantized"},
         "segmentation": {
             "categories",
             ("_metrics", "COCO-val2017-VOC-labels", "miou"),
             ("_metrics", "COCO-val2017-VOC-labels", "pixel_acc"),
-        },
-        "video": {"categories", ("_metrics", "Kinetics-400", "acc@1"), ("_metrics", "Kinetics-400", "acc@5")},
-        "optical_flow": set(),
+        }
+        | rich_metadata,
+        "video": {"categories", ("_metrics", "Kinetics-400", "acc@1"), ("_metrics", "Kinetics-400", "acc@5")}
+        | rich_metadata,
+        "optical_flow": rich_metadata,
     }
     model_name = model_fn.__name__
     module_name = model_fn.__module__.split(".")[-2]

torchvision/models/alexnet.py

@@ -67,6 +67,8 @@ class AlexNet_Weights(WeightsEnum):
                     "acc@5": 79.066,
                 }
             },
+            "_flops": 0.714188,
+            "_weight_size": 233.086501,
             "_docs": """
                 These weights reproduce closely the results of the paper using a simplified training recipe.
             """,

torchvision/models/convnext.py

@@ -219,6 +219,8 @@ class ConvNeXt_Tiny_Weights(WeightsEnum):
                     "acc@5": 96.146,
                 }
             },
+            "_flops": 4.455531,
+            "_weight_size": 109.118672,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -237,6 +239,8 @@ class ConvNeXt_Small_Weights(WeightsEnum):
                     "acc@5": 96.650,
                 }
             },
+            "_flops": 8.683712,
+            "_weight_size": 191.702775,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -255,6 +259,8 @@ class ConvNeXt_Base_Weights(WeightsEnum):
                     "acc@5": 96.870,
                 }
             },
+            "_flops": 15.354729,
+            "_weight_size": 338.064286,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -273,6 +279,8 @@ class ConvNeXt_Large_Weights(WeightsEnum):
                     "acc@5": 96.976,
                 }
             },
+            "_flops": 34.361434,
+            "_weight_size": 754.537187,
         },
     )
     DEFAULT = IMAGENET1K_V1

torchvision/models/densenet.py

@@ -15,7 +15,6 @@
 from ._meta import _IMAGENET_CATEGORIES
 from ._utils import _ovewrite_named_param, handle_legacy_interface
 
-
 __all__ = [
     "DenseNet",
     "DenseNet121_Weights",
@@ -278,6 +277,8 @@ class DenseNet121_Weights(WeightsEnum):
                     "acc@5": 91.972,
                 }
             },
+            "_flops": 2.834162,
+            "_weight_size": 30.844645,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -296,6 +297,8 @@ class DenseNet161_Weights(WeightsEnum):
                     "acc@5": 93.560,
                 }
             },
+            "_flops": 7.727907,
+            "_weight_size": 110.369482,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -314,6 +317,8 @@ class DenseNet169_Weights(WeightsEnum):
                     "acc@5": 92.806,
                 }
             },
+            "_flops": 3.359843,
+            "_weight_size": 54.708029,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -332,6 +337,8 @@ class DenseNet201_Weights(WeightsEnum):
                     "acc@5": 93.370,
                 }
             },
+            "_flops": 4.291366,
+            "_weight_size": 77.373247,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -444,7 +451,6 @@ def densenet201(*, weights: Optional[DenseNet201_Weights] = None, progress: bool
 # The dictionary below is internal implementation detail and will be removed in v0.15
 from ._utils import _ModelURLs
 
-
 model_urls = _ModelURLs(
     {
         "densenet121": DenseNet121_Weights.IMAGENET1K_V1.url,

torchvision/models/detection/faster_rcnn.py

@@ -388,6 +388,8 @@ class FasterRCNN_ResNet50_FPN_Weights(WeightsEnum):
                     "box_map": 37.0,
                 }
             },
+            "_flops": 134.379721,
+            "_weight_size": 159.743153,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )
@@ -407,6 +409,8 @@ class FasterRCNN_ResNet50_FPN_V2_Weights(WeightsEnum):
                     "box_map": 46.7,
                 }
             },
+            "_flops": 280.370729,
+            "_weight_size": 167.104394,
             "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
         },
     )
@@ -426,6 +430,8 @@ class FasterRCNN_MobileNet_V3_Large_FPN_Weights(WeightsEnum):
                     "box_map": 32.8,
                 }
             },
+            "_flops": 4.493592,
+            "_weight_size": 74.238593,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )
@@ -445,6 +451,8 @@ class FasterRCNN_MobileNet_V3_Large_320_FPN_Weights(WeightsEnum):
                     "box_map": 22.8,
                 }
             },
+            "_flops": 0.718998,
+            "_weight_size": 74.238593,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )

torchvision/models/detection/fcos.py

@@ -662,6 +662,8 @@ class FCOS_ResNet50_FPN_Weights(WeightsEnum):
                     "box_map": 39.2,
                 }
             },
+            "_flops": 128.207053,
+            "_weight_size": 123.607730,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )

torchvision/models/detection/keypoint_rcnn.py

@@ -328,6 +328,8 @@ class KeypointRCNN_ResNet50_FPN_Weights(WeightsEnum):
                     "kp_map": 61.1,
                 }
             },
+            "_flops": 133.924041,
+            "_weight_size": 226.053994,
             "_docs": """
                 These weights were produced by following a similar training recipe as on the paper but use a checkpoint
                 from an early epoch.
@@ -347,6 +349,8 @@ class KeypointRCNN_ResNet50_FPN_Weights(WeightsEnum):
                     "kp_map": 65.0,
                 }
             },
+            "_flops": 137.419502,
+            "_weight_size": 226.053994,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )

torchvision/models/detection/mask_rcnn.py

@@ -370,6 +370,8 @@ class MaskRCNN_ResNet50_FPN_Weights(WeightsEnum):
                     "mask_map": 34.6,
                 }
             },
+            "_flops": 134.379721,
+            "_weight_size": 169.839934,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )
@@ -390,6 +392,8 @@ class MaskRCNN_ResNet50_FPN_V2_Weights(WeightsEnum):
                     "mask_map": 41.8,
                 }
             },
+            "_flops": 333.577360,
+            "_weight_size": 177.219453,
             "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
         },
     )

torchvision/models/detection/retinanet.py

@@ -690,6 +690,8 @@ class RetinaNet_ResNet50_FPN_Weights(WeightsEnum):
                     "box_map": 36.4,
                 }
             },
+            "_flops": 151.540437,
+            "_weight_size": 130.267216,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )
@@ -709,6 +711,8 @@ class RetinaNet_ResNet50_FPN_V2_Weights(WeightsEnum):
                     "box_map": 41.5,
                 }
             },
+            "_flops": 152.238199,
+            "_weight_size": 146.037091,
             "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
         },
     )

torchvision/models/detection/ssd.py

@@ -39,6 +39,8 @@ class SSD300_VGG16_Weights(WeightsEnum):
                     "box_map": 25.1,
                 }
             },
+            "_flops": 34.858153,
+            "_weight_size": 135.988447,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )

torchvision/models/detection/ssdlite.py

@@ -198,6 +198,8 @@ class SSDLite320_MobileNet_V3_Large_Weights(WeightsEnum):
                     "box_map": 21.3,
                 }
             },
+            "_flops": 0.583172,
+            "_weight_size": 13.417583,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )