Skip to content

Commit f1587a2

Browse files
datumboxdatumbox
authored
Clean up purely informational fields from Weight Meta-data (#5852)
* Removing `task`, `architecture` and `quantization` * Fix mypy * Remove size field * Remove unused import. * Fix mypy * Remove size from schema list. * update todo * Simplify with assert * Adding min_size to all models. * Update RAFT min size to 128
1 parent 7998cdf commit f1587a2

34 files changed

+23
-160
lines changed

‎test/test_extended_models.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -79,12 +79,13 @@ def test_naming_conventions(model_fn):
7979
)
8080
@run_if_test_with_extended
8181
def test_schema_meta_validation(model_fn):
82-
classification_fields = ["size", "categories", "acc@1", "acc@5", "min_size"]
82+
# TODO: add list of permitted fields
83+
classification_fields = ["categories", "acc@1", "acc@5"]
8384
defaults = {
84-
"all": ["task", "architecture", "recipe", "num_params"],
85+
"all": ["recipe", "num_params", "min_size"],
8586
"models": classification_fields,
8687
"detection": ["categories", "map"],
87-
"quantization": classification_fields + ["backend", "quantization", "unquantized"],
88+
"quantization": classification_fields + ["backend", "unquantized"],
8889
"segmentation": ["categories", "mIoU", "acc"],
8990
"video": classification_fields,
9091
"optical_flow": [],

‎torchvision/models/alexnet.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,10 +57,7 @@ class AlexNet_Weights(WeightsEnum):
5757
url="https://download.pytorch.org/models/alexnet-owt-7be5be79.pth",
5858
transforms=partial(ImageClassification, crop_size=224),
5959
meta={
60-
"task": "image_classification",
61-
"architecture": "AlexNet",
6260
"num_params": 61100840,
63-
"size": (224, 224),
6461
"min_size": (63, 63),
6562
"categories": _IMAGENET_CATEGORIES,
6663
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#alexnet-and-vgg",

‎torchvision/models/convnext.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -204,9 +204,6 @@ def _convnext(
204204

205205

206206
_COMMON_META = {
207-
"task": "image_classification",
208-
"architecture": "ConvNeXt",
209-
"size": (224, 224),
210207
"min_size": (32, 32),
211208
"categories": _IMAGENET_CATEGORIES,
212209
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#convnext",

‎torchvision/models/densenet.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -266,9 +266,6 @@ def _densenet(
266266

267267

268268
_COMMON_META = {
269-
"task": "image_classification",
270-
"architecture": "DenseNet",
271-
"size": (224, 224),
272269
"min_size": (29, 29),
273270
"categories": _IMAGENET_CATEGORIES,
274271
"recipe": "https://github.com/pytorch/vision/pull/116",

‎torchvision/models/detection/faster_rcnn.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -370,9 +370,8 @@ def forward(self, x):
370370

371371

372372
_COMMON_META = {
373-
"task": "image_object_detection",
374-
"architecture": "FasterRCNN",
375373
"categories": _COCO_CATEGORIES,
374+
"min_size": (1, 1),
376375
}
377376

378377

‎torchvision/models/detection/fcos.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -651,10 +651,9 @@ class FCOS_ResNet50_FPN_Weights(WeightsEnum):
651651
url="https://download.pytorch.org/models/fcos_resnet50_fpn_coco-99b0c9b7.pth",
652652
transforms=ObjectDetection,
653653
meta={
654-
"task": "image_object_detection",
655-
"architecture": "FCOS",
656654
"num_params": 32269600,
657655
"categories": _COCO_CATEGORIES,
656+
"min_size": (1, 1),
658657
"recipe": "https://github.com/pytorch/vision/tree/main/references/detection#fcos-resnet-50-fpn",
659658
"map": 39.2,
660659
},

‎torchvision/models/detection/keypoint_rcnn.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -308,10 +308,9 @@ def forward(self, x):
308308

309309

310310
_COMMON_META = {
311-
"task": "image_object_detection",
312-
"architecture": "KeypointRCNN",
313311
"categories": _COCO_PERSON_CATEGORIES,
314312
"keypoint_names": _COCO_PERSON_KEYPOINT_NAMES,
313+
"min_size": (1, 1),
315314
}
316315

317316

‎torchvision/models/detection/mask_rcnn.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -351,9 +351,8 @@ def __init__(self, in_channels, dim_reduced, num_classes):
351351

352352

353353
_COMMON_META = {
354-
"task": "image_object_detection",
355-
"architecture": "MaskRCNN",
356354
"categories": _COCO_CATEGORIES,
355+
"min_size": (1, 1),
357356
}
358357

359358

‎torchvision/models/detection/retinanet.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -674,9 +674,8 @@ def forward(self, images, targets=None):
674674

675675

676676
_COMMON_META = {
677-
"task": "image_object_detection",
678-
"architecture": "RetinaNet",
679677
"categories": _COCO_CATEGORIES,
678+
"min_size": (1, 1),
680679
}
681680

682681

‎torchvision/models/detection/ssd.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,11 +30,9 @@ class SSD300_VGG16_Weights(WeightsEnum):
3030
url="https://download.pytorch.org/models/ssd300_vgg16_coco-b556d3b4.pth",
3131
transforms=ObjectDetection,
3232
meta={
33-
"task": "image_object_detection",
34-
"architecture": "SSD",
3533
"num_params": 35641826,
36-
"size": (300, 300),
3734
"categories": _COCO_CATEGORIES,
35+
"min_size": (1, 1),
3836
"recipe": "https://github.com/pytorch/vision/tree/main/references/detection#ssd300-vgg16",
3937
"map": 25.1,
4038
},

‎torchvision/models/detection/ssdlite.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -189,11 +189,9 @@ class SSDLite320_MobileNet_V3_Large_Weights(WeightsEnum):
189189
url="https://download.pytorch.org/models/ssdlite320_mobilenet_v3_large_coco-a79551df.pth",
190190
transforms=ObjectDetection,
191191
meta={
192-
"task": "image_object_detection",
193-
"architecture": "SSDLite",
194192
"num_params": 3440060,
195-
"size": (320, 320),
196193
"categories": _COCO_CATEGORIES,
194+
"min_size": (1, 1),
197195
"recipe": "https://github.com/pytorch/vision/tree/main/references/detection#ssdlite320-mobilenetv3-large",
198196
"map": 21.3,
199197
},

‎torchvision/models/efficientnet.py

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import warnings
44
from dataclasses import dataclass
55
from functools import partial
6-
from typing import Any, Callable, Optional, List, Sequence, Tuple, Union
6+
from typing import Any, Callable, Dict, Optional, List, Sequence, Tuple, Union
77

88
import torch
99
from torch import nn, Tensor
@@ -429,23 +429,20 @@ def _efficientnet_conf(
429429
return inverted_residual_setting, last_channel
430430

431431

432-
_COMMON_META = {
433-
"task": "image_classification",
432+
_COMMON_META: Dict[str, Any] = {
434433
"categories": _IMAGENET_CATEGORIES,
435434
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#efficientnet",
436435
}
437436

438437

439438
_COMMON_META_V1 = {
440439
**_COMMON_META,
441-
"architecture": "EfficientNet",
442440
"min_size": (1, 1),
443441
}
444442

445443

446444
_COMMON_META_V2 = {
447445
**_COMMON_META,
448-
"architecture": "EfficientNetV2",
449446
"min_size": (33, 33),
450447
}
451448

@@ -459,7 +456,6 @@ class EfficientNet_B0_Weights(WeightsEnum):
459456
meta={
460457
**_COMMON_META_V1,
461458
"num_params": 5288548,
462-
"size": (224, 224),
463459
"acc@1": 77.692,
464460
"acc@5": 93.532,
465461
},
@@ -476,7 +472,6 @@ class EfficientNet_B1_Weights(WeightsEnum):
476472
meta={
477473
**_COMMON_META_V1,
478474
"num_params": 7794184,
479-
"size": (240, 240),
480475
"acc@1": 78.642,
481476
"acc@5": 94.186,
482477
},
@@ -490,7 +485,6 @@ class EfficientNet_B1_Weights(WeightsEnum):
490485
**_COMMON_META_V1,
491486
"num_params": 7794184,
492487
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-lr-wd-crop-tuning",
493-
"size": (240, 240),
494488
"acc@1": 79.838,
495489
"acc@5": 94.934,
496490
},
@@ -507,7 +501,6 @@ class EfficientNet_B2_Weights(WeightsEnum):
507501
meta={
508502
**_COMMON_META_V1,
509503
"num_params": 9109994,
510-
"size": (288, 288),
511504
"acc@1": 80.608,
512505
"acc@5": 95.310,
513506
},
@@ -524,7 +517,6 @@ class EfficientNet_B3_Weights(WeightsEnum):
524517
meta={
525518
**_COMMON_META_V1,
526519
"num_params": 12233232,
527-
"size": (300, 300),
528520
"acc@1": 82.008,
529521
"acc@5": 96.054,
530522
},
@@ -541,7 +533,6 @@ class EfficientNet_B4_Weights(WeightsEnum):
541533
meta={
542534
**_COMMON_META_V1,
543535
"num_params": 19341616,
544-
"size": (380, 380),
545536
"acc@1": 83.384,
546537
"acc@5": 96.594,
547538
},
@@ -558,7 +549,6 @@ class EfficientNet_B5_Weights(WeightsEnum):
558549
meta={
559550
**_COMMON_META_V1,
560551
"num_params": 30389784,
561-
"size": (456, 456),
562552
"acc@1": 83.444,
563553
"acc@5": 96.628,
564554
},
@@ -575,7 +565,6 @@ class EfficientNet_B6_Weights(WeightsEnum):
575565
meta={
576566
**_COMMON_META_V1,
577567
"num_params": 43040704,
578-
"size": (528, 528),
579568
"acc@1": 84.008,
580569
"acc@5": 96.916,
581570
},
@@ -592,7 +581,6 @@ class EfficientNet_B7_Weights(WeightsEnum):
592581
meta={
593582
**_COMMON_META_V1,
594583
"num_params": 66347960,
595-
"size": (600, 600),
596584
"acc@1": 84.122,
597585
"acc@5": 96.908,
598586
},
@@ -612,7 +600,6 @@ class EfficientNet_V2_S_Weights(WeightsEnum):
612600
meta={
613601
**_COMMON_META_V2,
614602
"num_params": 21458488,
615-
"size": (384, 384),
616603
"acc@1": 84.228,
617604
"acc@5": 96.878,
618605
},
@@ -632,7 +619,6 @@ class EfficientNet_V2_M_Weights(WeightsEnum):
632619
meta={
633620
**_COMMON_META_V2,
634621
"num_params": 54139356,
635-
"size": (480, 480),
636622
"acc@1": 85.112,
637623
"acc@5": 97.156,
638624
},
@@ -654,7 +640,6 @@ class EfficientNet_V2_L_Weights(WeightsEnum):
654640
meta={
655641
**_COMMON_META_V2,
656642
"num_params": 118515272,
657-
"size": (480, 480),
658643
"acc@1": 85.808,
659644
"acc@5": 97.788,
660645
},

‎torchvision/models/googlenet.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -280,10 +280,7 @@ class GoogLeNet_Weights(WeightsEnum):
280280
url="https://download.pytorch.org/models/googlenet-1378be20.pth",
281281
transforms=partial(ImageClassification, crop_size=224),
282282
meta={
283-
"task": "image_classification",
284-
"architecture": "GoogLeNet",
285283
"num_params": 6624904,
286-
"size": (224, 224),
287284
"min_size": (15, 15),
288285
"categories": _IMAGENET_CATEGORIES,
289286
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#googlenet",

‎torchvision/models/inception.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -412,10 +412,7 @@ class Inception_V3_Weights(WeightsEnum):
412412
url="https://download.pytorch.org/models/inception_v3_google-0cc3c7bd.pth",
413413
transforms=partial(ImageClassification, crop_size=299, resize_size=342),
414414
meta={
415-
"task": "image_classification",
416-
"architecture": "InceptionV3",
417415
"num_params": 27161264,
418-
"size": (299, 299),
419416
"min_size": (75, 75),
420417
"categories": _IMAGENET_CATEGORIES,
421418
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#inception-v3",

‎torchvision/models/mnasnet.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -212,9 +212,6 @@ def _load_from_state_dict(
212212

213213

214214
_COMMON_META = {
215-
"task": "image_classification",
216-
"architecture": "MNASNet",
217-
"size": (224, 224),
218215
"min_size": (1, 1),
219216
"categories": _IMAGENET_CATEGORIES,
220217
"recipe": "https://github.com/1e100/mnasnet_trainer",

‎torchvision/models/mobilenetv2.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -195,10 +195,7 @@ def forward(self, x: Tensor) -> Tensor:
195195

196196

197197
_COMMON_META = {
198-
"task": "image_classification",
199-
"architecture": "MobileNetV2",
200198
"num_params": 3504872,
201-
"size": (224, 224),
202199
"min_size": (1, 1),
203200
"categories": _IMAGENET_CATEGORIES,
204201
}

‎torchvision/models/mobilenetv3.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -304,9 +304,6 @@ def _mobilenet_v3(
304304

305305

306306
_COMMON_META = {
307-
"task": "image_classification",
308-
"architecture": "MobileNetV3",
309-
"size": (224, 224),
310307
"min_size": (1, 1),
311308
"categories": _IMAGENET_CATEGORIES,
312309
}

‎torchvision/models/optical_flow/raft.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -512,8 +512,7 @@ def forward(self, image1, image2, num_flow_updates: int = 12):
512512

513513

514514
_COMMON_META = {
515-
"task": "optical_flow",
516-
"architecture": "RAFT",
515+
"min_size": (128, 128),
517516
}
518517

519518

‎torchvision/models/quantization/googlenet.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -111,14 +111,10 @@ class GoogLeNet_QuantizedWeights(WeightsEnum):
111111
url="https://download.pytorch.org/models/quantized/googlenet_fbgemm-c00238cf.pth",
112112
transforms=partial(ImageClassification, crop_size=224),
113113
meta={
114-
"task": "image_classification",
115-
"architecture": "GoogLeNet",
116114
"num_params": 6624904,
117-
"size": (224, 224),
118115
"min_size": (15, 15),
119116
"categories": _IMAGENET_CATEGORIES,
120117
"backend": "fbgemm",
121-
"quantization": "Post Training Quantization",
122118
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-models",
123119
"unquantized": GoogLeNet_Weights.IMAGENET1K_V1,
124120
"acc@1": 69.826,

‎torchvision/models/quantization/inception.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -177,14 +177,10 @@ class Inception_V3_QuantizedWeights(WeightsEnum):
177177
url="https://download.pytorch.org/models/quantized/inception_v3_google_fbgemm-71447a44.pth",
178178
transforms=partial(ImageClassification, crop_size=299, resize_size=342),
179179
meta={
180-
"task": "image_classification",
181-
"architecture": "InceptionV3",
182180
"num_params": 27161264,
183-
"size": (299, 299),
184181
"min_size": (75, 75),
185182
"categories": _IMAGENET_CATEGORIES,
186183
"backend": "fbgemm",
187-
"quantization": "Post Training Quantization",
188184
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-models",
189185
"unquantized": Inception_V3_Weights.IMAGENET1K_V1,
190186
"acc@1": 77.176,

‎torchvision/models/quantization/mobilenetv2.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -69,14 +69,10 @@ class MobileNet_V2_QuantizedWeights(WeightsEnum):
6969
url="https://download.pytorch.org/models/quantized/mobilenet_v2_qnnpack_37f702c5.pth",
7070
transforms=partial(ImageClassification, crop_size=224),
7171
meta={
72-
"task": "image_classification",
73-
"architecture": "MobileNetV2",
7472
"num_params": 3504872,
75-
"size": (224, 224),
7673
"min_size": (1, 1),
7774
"categories": _IMAGENET_CATEGORIES,
7875
"backend": "qnnpack",
79-
"quantization": "Quantization Aware Training",
8076
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#qat-mobilenetv2",
8177
"unquantized": MobileNet_V2_Weights.IMAGENET1K_V1,
8278
"acc@1": 71.658,

‎torchvision/models/quantization/mobilenetv3.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -159,14 +159,10 @@ class MobileNet_V3_Large_QuantizedWeights(WeightsEnum):
159159
url="https://download.pytorch.org/models/quantized/mobilenet_v3_large_qnnpack-5bcacf28.pth",
160160
transforms=partial(ImageClassification, crop_size=224),
161161
meta={
162-
"task": "image_classification",
163-
"architecture": "MobileNetV3",
164162
"num_params": 5483032,
165-
"size": (224, 224),
166163
"min_size": (1, 1),
167164
"categories": _IMAGENET_CATEGORIES,
168165
"backend": "qnnpack",
169-
"quantization": "Quantization Aware Training",
170166
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#qat-mobilenetv3",
171167
"unquantized": MobileNet_V3_Large_Weights.IMAGENET1K_V1,
172168
"acc@1": 73.004,

0 commit comments

Comments
 (0)