Cleanup weight docs (#7074)

* _weight_size -> _file_size

* Better formatting of individual weights tables

* Remove file size from main tables to avoid confusion with weight size (as in RAM)

* Remove unnecessary (file size) suffix

* Fix CI error?

* Formatting

Co-authored-by: Philip Meier <[email protected]>

Author: NicolasHug

docs/source/conf.py

@@ -363,12 +363,11 @@ def inject_weight_metadata(app, what, name, obj, options, lines):
                     v_sample = ", ".join(v[:max_visible])
                     v = f"{v_sample}, ... ({len(v)-max_visible} omitted)" if len(v) > max_visible else v_sample
                 elif k == "_ops":
-                    if obj.__name__.endswith("_QuantizedWeights"):
-                        v = f"{v} giga instructions per sec"
-                    else:
-                        v = f"{v} giga floating-point operations per sec"
-                elif k == "_weight_size":
-                    v = f"{v} MB (file size)"
+                    v = f"{v:.2f}"
+                    k = "GIPS" if obj.__name__.endswith("_QuantizedWeights") else "GFLOPS"
+                elif k == "_file_size":
+                    k = "File size"
+                    v = f"{v:.1f} MB"
 
                 table.append((str(k), str(v)))
             table = tabulate(table, tablefmt="rst")
@@ -396,9 +395,7 @@ def generate_weights_table(module, table_name, metrics, dataset, include_pattern
     ops_name = "GIPS" if "QuantizedWeights" in weights_endswith else "GFLOPS"
 
     metrics_keys, metrics_names = zip(*metrics)
-    column_names = (
-        ["Weight"] + list(metrics_names) + ["Params"] + [ops_name, "Size (MB)", "Recipe"]
-    )  # Final column order
+    column_names = ["Weight"] + list(metrics_names) + ["Params"] + [ops_name, "Recipe"]  # Final column order
     column_names = [f"**{name}**" for name in column_names]  # Add bold
 
     content = []
@@ -407,14 +404,13 @@ def generate_weights_table(module, table_name, metrics, dataset, include_pattern
             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"{w.meta['_ops']:.3f}",
-            f"{round(w.meta['_weight_size'], 1):.1f}",
+            f"{w.meta['_ops']:.2f}",
             f"`link <{w.meta['recipe']}>`__",
         ]
 
         content.append(row)
 
-    column_widths = ["110"] + ["18"] * len(metrics_names) + ["18"] * 3 + ["10"]
+    column_widths = ["110"] + ["18"] * len(metrics_names) + ["18"] * 2 + ["10"]
     widths_table = " ".join(column_widths)
 
     table = tabulate(content, headers=column_names, tablefmt="rst")

test/common_extended_utils.py

@@ -296,7 +296,7 @@ def get_ops(model: torch.nn.Module, weight: Weights, height=512, width=512):
     return round(flops, 3)
 
 
-def get_weight_size_mb(weight):
+def get_file_size_mb(weight):
     weights_path = os.path.join(os.getenv("HOME"), ".cache/torch/hub/checkpoints", weight.url.split("/")[-1])
     weights_size_mb = os.path.getsize(weights_path) / 1024 / 1024
 

test/test_extended_models.py

@@ -4,7 +4,7 @@
 import pytest
 import test_models as TM
 import torch
-from common_extended_utils import get_ops, get_weight_size_mb
+from common_extended_utils import get_file_size_mb, get_ops
 from torchvision import models
 from torchvision.models._api import get_model_weights, Weights, WeightsEnum
 from torchvision.models._utils import handle_legacy_interface
@@ -172,12 +172,12 @@ def test_schema_meta_validation(model_fn):
         "unquantized",
         "_docs",
         "_ops",
-        "_weight_size",
+        "_file_size",
     }
     # mandatory fields for each computer vision task
     classification_fields = {"categories", ("_metrics", "ImageNet-1K", "acc@1"), ("_metrics", "ImageNet-1K", "acc@5")}
     defaults = {
-        "all": {"_metrics", "min_size", "num_params", "recipe", "_docs", "_weight_size", "_ops"},
+        "all": {"_metrics", "min_size", "num_params", "recipe", "_docs", "_file_size", "_ops"},
         "models": classification_fields,
         "detection": {"categories", ("_metrics", "COCO-val2017", "box_map")},
         "quantization": classification_fields | {"backend", "unquantized"},
@@ -245,8 +245,8 @@ def test_schema_meta_validation(model_fn):
         if not w.name.isupper():
             bad_names.append(w)
 
-        if get_weight_size_mb(w) != w.meta.get("_weight_size"):
-            incorrect_meta.append((w, "_weight_size"))
+        if get_file_size_mb(w) != w.meta.get("_file_size"):
+            incorrect_meta.append((w, "_file_size"))
 
     assert not problematic_weights
     assert not incorrect_meta

torchvision/models/alexnet.py

@@ -68,7 +68,7 @@ class AlexNet_Weights(WeightsEnum):
                 }
             },
             "_ops": 0.714,
-            "_weight_size": 233.087,
+            "_file_size": 233.087,
             "_docs": """
                 These weights reproduce closely the results of the paper using a simplified training recipe.
             """,

torchvision/models/convnext.py

@@ -220,7 +220,7 @@ class ConvNeXt_Tiny_Weights(WeightsEnum):
                 }
             },
             "_ops": 4.456,
-            "_weight_size": 109.119,
+            "_file_size": 109.119,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -240,7 +240,7 @@ class ConvNeXt_Small_Weights(WeightsEnum):
                 }
             },
             "_ops": 8.684,
-            "_weight_size": 191.703,
+            "_file_size": 191.703,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -260,7 +260,7 @@ class ConvNeXt_Base_Weights(WeightsEnum):
                 }
             },
             "_ops": 15.355,
-            "_weight_size": 338.064,
+            "_file_size": 338.064,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -280,7 +280,7 @@ class ConvNeXt_Large_Weights(WeightsEnum):
                 }
             },
             "_ops": 34.361,
-            "_weight_size": 754.537,
+            "_file_size": 754.537,
         },
     )
     DEFAULT = IMAGENET1K_V1

torchvision/models/densenet.py

@@ -278,7 +278,7 @@ class DenseNet121_Weights(WeightsEnum):
                 }
             },
             "_ops": 2.834,
-            "_weight_size": 30.845,
+            "_file_size": 30.845,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -298,7 +298,7 @@ class DenseNet161_Weights(WeightsEnum):
                 }
             },
             "_ops": 7.728,
-            "_weight_size": 110.369,
+            "_file_size": 110.369,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -318,7 +318,7 @@ class DenseNet169_Weights(WeightsEnum):
                 }
             },
             "_ops": 3.36,
-            "_weight_size": 54.708,
+            "_file_size": 54.708,
         },
     )
     DEFAULT = IMAGENET1K_V1
@@ -338,7 +338,7 @@ class DenseNet201_Weights(WeightsEnum):
                 }
             },
             "_ops": 4.291,
-            "_weight_size": 77.373,
+            "_file_size": 77.373,
         },
     )
     DEFAULT = IMAGENET1K_V1

torchvision/models/detection/faster_rcnn.py

@@ -389,7 +389,7 @@ class FasterRCNN_ResNet50_FPN_Weights(WeightsEnum):
                 }
             },
             "_ops": 134.38,
-            "_weight_size": 159.743,
+            "_file_size": 159.743,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )
@@ -410,7 +410,7 @@ class FasterRCNN_ResNet50_FPN_V2_Weights(WeightsEnum):
                 }
             },
             "_ops": 280.371,
-            "_weight_size": 167.104,
+            "_file_size": 167.104,
             "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
         },
     )
@@ -431,7 +431,7 @@ class FasterRCNN_MobileNet_V3_Large_FPN_Weights(WeightsEnum):
                 }
             },
             "_ops": 4.494,
-            "_weight_size": 74.239,
+            "_file_size": 74.239,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )
@@ -452,7 +452,7 @@ class FasterRCNN_MobileNet_V3_Large_320_FPN_Weights(WeightsEnum):
                 }
             },
             "_ops": 0.719,
-            "_weight_size": 74.239,
+            "_file_size": 74.239,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )

torchvision/models/detection/fcos.py

@@ -663,7 +663,7 @@ class FCOS_ResNet50_FPN_Weights(WeightsEnum):
                 }
             },
             "_ops": 128.207,
-            "_weight_size": 123.608,
+            "_file_size": 123.608,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )

torchvision/models/detection/keypoint_rcnn.py

@@ -329,7 +329,7 @@ class KeypointRCNN_ResNet50_FPN_Weights(WeightsEnum):
                 }
             },
             "_ops": 133.924,
-            "_weight_size": 226.054,
+            "_file_size": 226.054,
             "_docs": """
                 These weights were produced by following a similar training recipe as on the paper but use a checkpoint
                 from an early epoch.
@@ -350,7 +350,7 @@ class KeypointRCNN_ResNet50_FPN_Weights(WeightsEnum):
                 }
             },
             "_ops": 137.42,
-            "_weight_size": 226.054,
+            "_file_size": 226.054,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )

torchvision/models/detection/mask_rcnn.py

@@ -371,7 +371,7 @@ class MaskRCNN_ResNet50_FPN_Weights(WeightsEnum):
                 }
             },
             "_ops": 134.38,
-            "_weight_size": 169.84,
+            "_file_size": 169.84,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )
@@ -393,7 +393,7 @@ class MaskRCNN_ResNet50_FPN_V2_Weights(WeightsEnum):
                 }
             },
             "_ops": 333.577,
-            "_weight_size": 177.219,
+            "_file_size": 177.219,
             "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
         },
     )