Skip to content

Commit ed2a0ad

Browse files
NicolasHugpmeier
andauthored
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]>
1 parent 90cfb10 commit ed2a0ad

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

41 files changed

+180
-184
lines changed

docs/source/conf.py

+8-12
Original file line numberDiff line numberDiff line change
@@ -363,12 +363,11 @@ def inject_weight_metadata(app, what, name, obj, options, lines):
363363
v_sample = ", ".join(v[:max_visible])
364364
v = f"{v_sample}, ... ({len(v)-max_visible} omitted)" if len(v) > max_visible else v_sample
365365
elif k == "_ops":
366-
if obj.__name__.endswith("_QuantizedWeights"):
367-
v = f"{v} giga instructions per sec"
368-
else:
369-
v = f"{v} giga floating-point operations per sec"
370-
elif k == "_weight_size":
371-
v = f"{v} MB (file size)"
366+
v = f"{v:.2f}"
367+
k = "GIPS" if obj.__name__.endswith("_QuantizedWeights") else "GFLOPS"
368+
elif k == "_file_size":
369+
k = "File size"
370+
v = f"{v:.1f} MB"
372371

373372
table.append((str(k), str(v)))
374373
table = tabulate(table, tablefmt="rst")
@@ -396,9 +395,7 @@ def generate_weights_table(module, table_name, metrics, dataset, include_pattern
396395
ops_name = "GIPS" if "QuantizedWeights" in weights_endswith else "GFLOPS"
397396

398397
metrics_keys, metrics_names = zip(*metrics)
399-
column_names = (
400-
["Weight"] + list(metrics_names) + ["Params"] + [ops_name, "Size (MB)", "Recipe"]
401-
) # Final column order
398+
column_names = ["Weight"] + list(metrics_names) + ["Params"] + [ops_name, "Recipe"] # Final column order
402399
column_names = [f"**{name}**" for name in column_names] # Add bold
403400

404401
content = []
@@ -407,14 +404,13 @@ def generate_weights_table(module, table_name, metrics, dataset, include_pattern
407404
f":class:`{w} <{type(w).__name__}>`",
408405
*(w.meta["_metrics"][dataset][metric] for metric in metrics_keys),
409406
f"{w.meta['num_params']/1e6:.1f}M",
410-
f"{w.meta['_ops']:.3f}",
411-
f"{round(w.meta['_weight_size'], 1):.1f}",
407+
f"{w.meta['_ops']:.2f}",
412408
f"`link <{w.meta['recipe']}>`__",
413409
]
414410

415411
content.append(row)
416412

417-
column_widths = ["110"] + ["18"] * len(metrics_names) + ["18"] * 3 + ["10"]
413+
column_widths = ["110"] + ["18"] * len(metrics_names) + ["18"] * 2 + ["10"]
418414
widths_table = " ".join(column_widths)
419415

420416
table = tabulate(content, headers=column_names, tablefmt="rst")

test/common_extended_utils.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -296,7 +296,7 @@ def get_ops(model: torch.nn.Module, weight: Weights, height=512, width=512):
296296
return round(flops, 3)
297297

298298

299-
def get_weight_size_mb(weight):
299+
def get_file_size_mb(weight):
300300
weights_path = os.path.join(os.getenv("HOME"), ".cache/torch/hub/checkpoints", weight.url.split("/")[-1])
301301
weights_size_mb = os.path.getsize(weights_path) / 1024 / 1024
302302

test/test_extended_models.py

+5-5
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import pytest
55
import test_models as TM
66
import torch
7-
from common_extended_utils import get_ops, get_weight_size_mb
7+
from common_extended_utils import get_file_size_mb, get_ops
88
from torchvision import models
99
from torchvision.models._api import get_model_weights, Weights, WeightsEnum
1010
from torchvision.models._utils import handle_legacy_interface
@@ -172,12 +172,12 @@ def test_schema_meta_validation(model_fn):
172172
"unquantized",
173173
"_docs",
174174
"_ops",
175-
"_weight_size",
175+
"_file_size",
176176
}
177177
# mandatory fields for each computer vision task
178178
classification_fields = {"categories", ("_metrics", "ImageNet-1K", "acc@1"), ("_metrics", "ImageNet-1K", "acc@5")}
179179
defaults = {
180-
"all": {"_metrics", "min_size", "num_params", "recipe", "_docs", "_weight_size", "_ops"},
180+
"all": {"_metrics", "min_size", "num_params", "recipe", "_docs", "_file_size", "_ops"},
181181
"models": classification_fields,
182182
"detection": {"categories", ("_metrics", "COCO-val2017", "box_map")},
183183
"quantization": classification_fields | {"backend", "unquantized"},
@@ -245,8 +245,8 @@ def test_schema_meta_validation(model_fn):
245245
if not w.name.isupper():
246246
bad_names.append(w)
247247

248-
if get_weight_size_mb(w) != w.meta.get("_weight_size"):
249-
incorrect_meta.append((w, "_weight_size"))
248+
if get_file_size_mb(w) != w.meta.get("_file_size"):
249+
incorrect_meta.append((w, "_file_size"))
250250

251251
assert not problematic_weights
252252
assert not incorrect_meta

torchvision/models/alexnet.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ class AlexNet_Weights(WeightsEnum):
6868
}
6969
},
7070
"_ops": 0.714,
71-
"_weight_size": 233.087,
71+
"_file_size": 233.087,
7272
"_docs": """
7373
These weights reproduce closely the results of the paper using a simplified training recipe.
7474
""",

torchvision/models/convnext.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -220,7 +220,7 @@ class ConvNeXt_Tiny_Weights(WeightsEnum):
220220
}
221221
},
222222
"_ops": 4.456,
223-
"_weight_size": 109.119,
223+
"_file_size": 109.119,
224224
},
225225
)
226226
DEFAULT = IMAGENET1K_V1
@@ -240,7 +240,7 @@ class ConvNeXt_Small_Weights(WeightsEnum):
240240
}
241241
},
242242
"_ops": 8.684,
243-
"_weight_size": 191.703,
243+
"_file_size": 191.703,
244244
},
245245
)
246246
DEFAULT = IMAGENET1K_V1
@@ -260,7 +260,7 @@ class ConvNeXt_Base_Weights(WeightsEnum):
260260
}
261261
},
262262
"_ops": 15.355,
263-
"_weight_size": 338.064,
263+
"_file_size": 338.064,
264264
},
265265
)
266266
DEFAULT = IMAGENET1K_V1
@@ -280,7 +280,7 @@ class ConvNeXt_Large_Weights(WeightsEnum):
280280
}
281281
},
282282
"_ops": 34.361,
283-
"_weight_size": 754.537,
283+
"_file_size": 754.537,
284284
},
285285
)
286286
DEFAULT = IMAGENET1K_V1

torchvision/models/densenet.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,7 @@ class DenseNet121_Weights(WeightsEnum):
278278
}
279279
},
280280
"_ops": 2.834,
281-
"_weight_size": 30.845,
281+
"_file_size": 30.845,
282282
},
283283
)
284284
DEFAULT = IMAGENET1K_V1
@@ -298,7 +298,7 @@ class DenseNet161_Weights(WeightsEnum):
298298
}
299299
},
300300
"_ops": 7.728,
301-
"_weight_size": 110.369,
301+
"_file_size": 110.369,
302302
},
303303
)
304304
DEFAULT = IMAGENET1K_V1
@@ -318,7 +318,7 @@ class DenseNet169_Weights(WeightsEnum):
318318
}
319319
},
320320
"_ops": 3.36,
321-
"_weight_size": 54.708,
321+
"_file_size": 54.708,
322322
},
323323
)
324324
DEFAULT = IMAGENET1K_V1
@@ -338,7 +338,7 @@ class DenseNet201_Weights(WeightsEnum):
338338
}
339339
},
340340
"_ops": 4.291,
341-
"_weight_size": 77.373,
341+
"_file_size": 77.373,
342342
},
343343
)
344344
DEFAULT = IMAGENET1K_V1

torchvision/models/detection/faster_rcnn.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -389,7 +389,7 @@ class FasterRCNN_ResNet50_FPN_Weights(WeightsEnum):
389389
}
390390
},
391391
"_ops": 134.38,
392-
"_weight_size": 159.743,
392+
"_file_size": 159.743,
393393
"_docs": """These weights were produced by following a similar training recipe as on the paper.""",
394394
},
395395
)
@@ -410,7 +410,7 @@ class FasterRCNN_ResNet50_FPN_V2_Weights(WeightsEnum):
410410
}
411411
},
412412
"_ops": 280.371,
413-
"_weight_size": 167.104,
413+
"_file_size": 167.104,
414414
"_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
415415
},
416416
)
@@ -431,7 +431,7 @@ class FasterRCNN_MobileNet_V3_Large_FPN_Weights(WeightsEnum):
431431
}
432432
},
433433
"_ops": 4.494,
434-
"_weight_size": 74.239,
434+
"_file_size": 74.239,
435435
"_docs": """These weights were produced by following a similar training recipe as on the paper.""",
436436
},
437437
)
@@ -452,7 +452,7 @@ class FasterRCNN_MobileNet_V3_Large_320_FPN_Weights(WeightsEnum):
452452
}
453453
},
454454
"_ops": 0.719,
455-
"_weight_size": 74.239,
455+
"_file_size": 74.239,
456456
"_docs": """These weights were produced by following a similar training recipe as on the paper.""",
457457
},
458458
)

torchvision/models/detection/fcos.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -663,7 +663,7 @@ class FCOS_ResNet50_FPN_Weights(WeightsEnum):
663663
}
664664
},
665665
"_ops": 128.207,
666-
"_weight_size": 123.608,
666+
"_file_size": 123.608,
667667
"_docs": """These weights were produced by following a similar training recipe as on the paper.""",
668668
},
669669
)

torchvision/models/detection/keypoint_rcnn.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -329,7 +329,7 @@ class KeypointRCNN_ResNet50_FPN_Weights(WeightsEnum):
329329
}
330330
},
331331
"_ops": 133.924,
332-
"_weight_size": 226.054,
332+
"_file_size": 226.054,
333333
"_docs": """
334334
These weights were produced by following a similar training recipe as on the paper but use a checkpoint
335335
from an early epoch.
@@ -350,7 +350,7 @@ class KeypointRCNN_ResNet50_FPN_Weights(WeightsEnum):
350350
}
351351
},
352352
"_ops": 137.42,
353-
"_weight_size": 226.054,
353+
"_file_size": 226.054,
354354
"_docs": """These weights were produced by following a similar training recipe as on the paper.""",
355355
},
356356
)

torchvision/models/detection/mask_rcnn.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -371,7 +371,7 @@ class MaskRCNN_ResNet50_FPN_Weights(WeightsEnum):
371371
}
372372
},
373373
"_ops": 134.38,
374-
"_weight_size": 169.84,
374+
"_file_size": 169.84,
375375
"_docs": """These weights were produced by following a similar training recipe as on the paper.""",
376376
},
377377
)
@@ -393,7 +393,7 @@ class MaskRCNN_ResNet50_FPN_V2_Weights(WeightsEnum):
393393
}
394394
},
395395
"_ops": 333.577,
396-
"_weight_size": 177.219,
396+
"_file_size": 177.219,
397397
"_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
398398
},
399399
)

torchvision/models/detection/retinanet.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -691,7 +691,7 @@ class RetinaNet_ResNet50_FPN_Weights(WeightsEnum):
691691
}
692692
},
693693
"_ops": 151.54,
694-
"_weight_size": 130.267,
694+
"_file_size": 130.267,
695695
"_docs": """These weights were produced by following a similar training recipe as on the paper.""",
696696
},
697697
)
@@ -712,7 +712,7 @@ class RetinaNet_ResNet50_FPN_V2_Weights(WeightsEnum):
712712
}
713713
},
714714
"_ops": 152.238,
715-
"_weight_size": 146.037,
715+
"_file_size": 146.037,
716716
"_docs": """These weights were produced using an enhanced training recipe to boost the model accuracy.""",
717717
},
718718
)

torchvision/models/detection/ssd.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ class SSD300_VGG16_Weights(WeightsEnum):
4040
}
4141
},
4242
"_ops": 34.858,
43-
"_weight_size": 135.988,
43+
"_file_size": 135.988,
4444
"_docs": """These weights were produced by following a similar training recipe as on the paper.""",
4545
},
4646
)

torchvision/models/detection/ssdlite.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -199,7 +199,7 @@ class SSDLite320_MobileNet_V3_Large_Weights(WeightsEnum):
199199
}
200200
},
201201
"_ops": 0.583,
202-
"_weight_size": 13.418,
202+
"_file_size": 13.418,
203203
"_docs": """These weights were produced by following a similar training recipe as on the paper.""",
204204
},
205205
)

torchvision/models/efficientnet.py

+12-12
Original file line numberDiff line numberDiff line change
@@ -465,7 +465,7 @@ class EfficientNet_B0_Weights(WeightsEnum):
465465
}
466466
},
467467
"_ops": 0.386,
468-
"_weight_size": 20.451,
468+
"_file_size": 20.451,
469469
"_docs": """These weights are ported from the original paper.""",
470470
},
471471
)
@@ -489,7 +489,7 @@ class EfficientNet_B1_Weights(WeightsEnum):
489489
}
490490
},
491491
"_ops": 0.687,
492-
"_weight_size": 30.134,
492+
"_file_size": 30.134,
493493
"_docs": """These weights are ported from the original paper.""",
494494
},
495495
)
@@ -509,7 +509,7 @@ class EfficientNet_B1_Weights(WeightsEnum):
509509
}
510510
},
511511
"_ops": 0.687,
512-
"_weight_size": 30.136,
512+
"_file_size": 30.136,
513513
"_docs": """
514514
These weights improve upon the results of the original paper by using a modified version of TorchVision's
515515
`new training recipe
@@ -537,7 +537,7 @@ class EfficientNet_B2_Weights(WeightsEnum):
537537
}
538538
},
539539
"_ops": 1.088,
540-
"_weight_size": 35.174,
540+
"_file_size": 35.174,
541541
"_docs": """These weights are ported from the original paper.""",
542542
},
543543
)
@@ -561,7 +561,7 @@ class EfficientNet_B3_Weights(WeightsEnum):
561561
}
562562
},
563563
"_ops": 1.827,
564-
"_weight_size": 47.184,
564+
"_file_size": 47.184,
565565
"_docs": """These weights are ported from the original paper.""",
566566
},
567567
)
@@ -585,7 +585,7 @@ class EfficientNet_B4_Weights(WeightsEnum):
585585
}
586586
},
587587
"_ops": 4.394,
588-
"_weight_size": 74.489,
588+
"_file_size": 74.489,
589589
"_docs": """These weights are ported from the original paper.""",
590590
},
591591
)
@@ -609,7 +609,7 @@ class EfficientNet_B5_Weights(WeightsEnum):
609609
}
610610
},
611611
"_ops": 10.266,
612-
"_weight_size": 116.864,
612+
"_file_size": 116.864,
613613
"_docs": """These weights are ported from the original paper.""",
614614
},
615615
)
@@ -633,7 +633,7 @@ class EfficientNet_B6_Weights(WeightsEnum):
633633
}
634634
},
635635
"_ops": 19.068,
636-
"_weight_size": 165.362,
636+
"_file_size": 165.362,
637637
"_docs": """These weights are ported from the original paper.""",
638638
},
639639
)
@@ -657,7 +657,7 @@ class EfficientNet_B7_Weights(WeightsEnum):
657657
}
658658
},
659659
"_ops": 37.746,
660-
"_weight_size": 254.675,
660+
"_file_size": 254.675,
661661
"_docs": """These weights are ported from the original paper.""",
662662
},
663663
)
@@ -683,7 +683,7 @@ class EfficientNet_V2_S_Weights(WeightsEnum):
683683
}
684684
},
685685
"_ops": 8.366,
686-
"_weight_size": 82.704,
686+
"_file_size": 82.704,
687687
"_docs": """
688688
These weights improve upon the results of the original paper by using a modified version of TorchVision's
689689
`new training recipe
@@ -713,7 +713,7 @@ class EfficientNet_V2_M_Weights(WeightsEnum):
713713
}
714714
},
715715
"_ops": 24.582,
716-
"_weight_size": 208.01,
716+
"_file_size": 208.01,
717717
"_docs": """
718718
These weights improve upon the results of the original paper by using a modified version of TorchVision's
719719
`new training recipe
@@ -746,7 +746,7 @@ class EfficientNet_V2_L_Weights(WeightsEnum):
746746
}
747747
},
748748
"_ops": 56.08,
749-
"_weight_size": 454.573,
749+
"_file_size": 454.573,
750750
"_docs": """These weights are ported from the original paper.""",
751751
},
752752
)

0 commit comments

Comments
 (0)