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 <github.pmeier@posteo.de>
pytorch · Jan 11, 2023 · ed2a0ad · ed2a0ad
1 parent 90cfb10
commit ed2a0ad
Show file tree

Hide file tree

Showing 41 changed files with 180 additions and 184 deletions.
diff --git a/docs/source/conf.py b/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")

diff --git a/test/common_extended_utils.py b/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
 

diff --git a/test/test_extended_models.py b/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

diff --git a/torchvision/models/alexnet.py b/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.
             """,

diff --git a/torchvision/models/convnext.py b/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

diff --git a/torchvision/models/densenet.py b/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

diff --git a/torchvision/models/detection/faster_rcnn.py b/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.""",
         },
     )

diff --git a/torchvision/models/detection/fcos.py b/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.""",
         },
     )

diff --git a/torchvision/models/detection/keypoint_rcnn.py b/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.""",
         },
     )

diff --git a/torchvision/models/detection/mask_rcnn.py b/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.""",
         },
     )

diff --git a/torchvision/models/detection/retinanet.py b/torchvision/models/detection/retinanet.py
@@ -691,7 +691,7 @@ class RetinaNet_ResNet50_FPN_Weights(WeightsEnum):
                 }
             },
             "_ops": 151.54,
-            "_weight_size": 130.267,
+            "_file_size": 130.267,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )
@@ -712,7 +712,7 @@ class RetinaNet_ResNet50_FPN_V2_Weights(WeightsEnum):
                 }
             },
             "_ops": 152.238,
-            "_weight_size": 146.037,
+            "_file_size": 146.037,
             "_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
         },
     )

diff --git a/torchvision/models/detection/ssd.py b/torchvision/models/detection/ssd.py
@@ -40,7 +40,7 @@ class SSD300_VGG16_Weights(WeightsEnum):
                 }
             },
             "_ops": 34.858,
-            "_weight_size": 135.988,
+            "_file_size": 135.988,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )

diff --git a/torchvision/models/detection/ssdlite.py b/torchvision/models/detection/ssdlite.py
@@ -199,7 +199,7 @@ class SSDLite320_MobileNet_V3_Large_Weights(WeightsEnum):
                 }
             },
             "_ops": 0.583,
-            "_weight_size": 13.418,
+            "_file_size": 13.418,
             "_docs": """These weights were produced by following a similar training recipe as on the paper.""",
         },
     )

diff --git a/torchvision/models/efficientnet.py b/torchvision/models/efficientnet.py
@@ -465,7 +465,7 @@ class EfficientNet_B0_Weights(WeightsEnum):
                 }
             },
             "_ops": 0.386,
-            "_weight_size": 20.451,
+            "_file_size": 20.451,
             "_docs": """These weights are ported from the original paper.""",
         },
     )
@@ -489,7 +489,7 @@ class EfficientNet_B1_Weights(WeightsEnum):
                 }
             },
             "_ops": 0.687,
-            "_weight_size": 30.134,
+            "_file_size": 30.134,
             "_docs": """These weights are ported from the original paper.""",
         },
     )
@@ -509,7 +509,7 @@ class EfficientNet_B1_Weights(WeightsEnum):
                 }
             },
             "_ops": 0.687,
-            "_weight_size": 30.136,
+            "_file_size": 30.136,
             "_docs": """
                 These weights improve upon the results of the original paper by using a modified version of TorchVision's
                 `new training recipe
@@ -537,7 +537,7 @@ class EfficientNet_B2_Weights(WeightsEnum):
                 }
             },
             "_ops": 1.088,
-            "_weight_size": 35.174,
+            "_file_size": 35.174,
             "_docs": """These weights are ported from the original paper.""",
         },
     )
@@ -561,7 +561,7 @@ class EfficientNet_B3_Weights(WeightsEnum):
                 }
             },
             "_ops": 1.827,
-            "_weight_size": 47.184,
+            "_file_size": 47.184,
             "_docs": """These weights are ported from the original paper.""",
         },
     )
@@ -585,7 +585,7 @@ class EfficientNet_B4_Weights(WeightsEnum):
                 }
             },
             "_ops": 4.394,
-            "_weight_size": 74.489,
+            "_file_size": 74.489,
             "_docs": """These weights are ported from the original paper.""",
         },
     )
@@ -609,7 +609,7 @@ class EfficientNet_B5_Weights(WeightsEnum):
                 }
             },
             "_ops": 10.266,
-            "_weight_size": 116.864,
+            "_file_size": 116.864,
             "_docs": """These weights are ported from the original paper.""",
         },
     )
@@ -633,7 +633,7 @@ class EfficientNet_B6_Weights(WeightsEnum):
                 }
             },
             "_ops": 19.068,
-            "_weight_size": 165.362,
+            "_file_size": 165.362,
             "_docs": """These weights are ported from the original paper.""",
         },
     )
@@ -657,7 +657,7 @@ class EfficientNet_B7_Weights(WeightsEnum):
                 }
             },
             "_ops": 37.746,
-            "_weight_size": 254.675,
+            "_file_size": 254.675,
             "_docs": """These weights are ported from the original paper.""",
         },
     )
@@ -683,7 +683,7 @@ class EfficientNet_V2_S_Weights(WeightsEnum):
                 }
             },
             "_ops": 8.366,
-            "_weight_size": 82.704,
+            "_file_size": 82.704,
             "_docs": """
                 These weights improve upon the results of the original paper by using a modified version of TorchVision's
                 `new training recipe
@@ -713,7 +713,7 @@ class EfficientNet_V2_M_Weights(WeightsEnum):
                 }
             },
             "_ops": 24.582,
-            "_weight_size": 208.01,
+            "_file_size": 208.01,
             "_docs": """
                 These weights improve upon the results of the original paper by using a modified version of TorchVision's
                 `new training recipe
@@ -746,7 +746,7 @@ class EfficientNet_V2_L_Weights(WeightsEnum):
                 }
             },
             "_ops": 56.08,
-            "_weight_size": 454.573,
+            "_file_size": 454.573,
             "_docs": """These weights are ported from the original paper.""",
         },
     )