Skip to content

GitLab

Unverified Commit f1587a20 authored by Vasilis Vryniotis's avatar Vasilis Vryniotis Committed by GitHub
Browse files

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
parent 7998cdfa
Hide whitespace changes
Inline Side-by-side
Showing with 14 additions and 70 deletions
test/test_extended_models.py
View file @ f1587a20
......@@ -79,12 +79,13 @@ def test_naming_conventions(model_fn):
)
@run_if_test_with_extended
def test_schema_meta_validation(model_fn):
classification_fields = ["size", "categories", "acc@1", "acc@5", "min_size"]
# TODO: add list of permitted fields
classification_fields = ["categories", "acc@1", "acc@5"]
defaults = {
"all": ["task", "architecture", "recipe", "num_params"],
"all": ["recipe", "num_params", "min_size"],
"models": classification_fields,
"detection": ["categories", "map"],
"quantization": classification_fields + ["backend", "quantization", "unquantized"],
"quantization": classification_fields + ["backend", "unquantized"],
"segmentation": ["categories", "mIoU", "acc"],
"video": classification_fields,
"optical_flow": [],
......
torchvision/models/alexnet.py
View file @ f1587a20
......@@ -57,10 +57,7 @@ class AlexNet_Weights(WeightsEnum):
url="https://download.pytorch.org/models/alexnet-owt-7be5be79.pth",
transforms=partial(ImageClassification, crop_size=224),
meta={
"task": "image_classification",
"architecture": "AlexNet",
"num_params": 61100840,
"size": (224, 224),
"min_size": (63, 63),
"categories": _IMAGENET_CATEGORIES,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#alexnet-and-vgg",
......
torchvision/models/convnext.py
View file @ f1587a20
......@@ -204,9 +204,6 @@ def _convnext(
_COMMON_META = {
"task": "image_classification",
"architecture": "ConvNeXt",
"size": (224, 224),
"min_size": (32, 32),
"categories": _IMAGENET_CATEGORIES,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#convnext",
......
torchvision/models/densenet.py
View file @ f1587a20
......@@ -266,9 +266,6 @@ def _densenet(
_COMMON_META = {
"task": "image_classification",
"architecture": "DenseNet",
"size": (224, 224),
"min_size": (29, 29),
"categories": _IMAGENET_CATEGORIES,
"recipe": "https://github.com/pytorch/vision/pull/116",
......
torchvision/models/detection/faster_rcnn.py
View file @ f1587a20
......@@ -370,9 +370,8 @@ class FastRCNNPredictor(nn.Module):
_COMMON_META = {
"task": "image_object_detection",
"architecture": "FasterRCNN",
"categories": _COCO_CATEGORIES,
"min_size": (1, 1),
}
......
torchvision/models/detection/fcos.py
View file @ f1587a20
......@@ -651,10 +651,9 @@ class FCOS_ResNet50_FPN_Weights(WeightsEnum):
url="https://download.pytorch.org/models/fcos_resnet50_fpn_coco-99b0c9b7.pth",
transforms=ObjectDetection,
meta={
"task": "image_object_detection",
"architecture": "FCOS",
"num_params": 32269600,
"categories": _COCO_CATEGORIES,
"min_size": (1, 1),
"recipe": "https://github.com/pytorch/vision/tree/main/references/detection#fcos-resnet-50-fpn",
"map": 39.2,
},
......
torchvision/models/detection/keypoint_rcnn.py
View file @ f1587a20
......@@ -308,10 +308,9 @@ class KeypointRCNNPredictor(nn.Module):
_COMMON_META = {
"task": "image_object_detection",
"architecture": "KeypointRCNN",
"categories": _COCO_PERSON_CATEGORIES,
"keypoint_names": _COCO_PERSON_KEYPOINT_NAMES,
"min_size": (1, 1),
}
......
torchvision/models/detection/mask_rcnn.py
View file @ f1587a20
......@@ -351,9 +351,8 @@ class MaskRCNNPredictor(nn.Sequential):
_COMMON_META = {
"task": "image_object_detection",
"architecture": "MaskRCNN",
"categories": _COCO_CATEGORIES,
"min_size": (1, 1),
}
......
torchvision/models/detection/retinanet.py
View file @ f1587a20
......@@ -674,9 +674,8 @@ class RetinaNet(nn.Module):
_COMMON_META = {
"task": "image_object_detection",
"architecture": "RetinaNet",
"categories": _COCO_CATEGORIES,
"min_size": (1, 1),
}
......
torchvision/models/detection/ssd.py
View file @ f1587a20
......@@ -30,11 +30,9 @@ class SSD300_VGG16_Weights(WeightsEnum):
url="https://download.pytorch.org/models/ssd300_vgg16_coco-b556d3b4.pth",
transforms=ObjectDetection,
meta={
"task": "image_object_detection",
"architecture": "SSD",
"num_params": 35641826,
"size": (300, 300),
"categories": _COCO_CATEGORIES,
"min_size": (1, 1),
"recipe": "https://github.com/pytorch/vision/tree/main/references/detection#ssd300-vgg16",
"map": 25.1,
},
......
torchvision/models/detection/ssdlite.py
View file @ f1587a20
......@@ -189,11 +189,9 @@ class SSDLite320_MobileNet_V3_Large_Weights(WeightsEnum):
url="https://download.pytorch.org/models/ssdlite320_mobilenet_v3_large_coco-a79551df.pth",
transforms=ObjectDetection,
meta={
"task": "image_object_detection",
"architecture": "SSDLite",
"num_params": 3440060,
"size": (320, 320),
"categories": _COCO_CATEGORIES,
"min_size": (1, 1),
"recipe": "https://github.com/pytorch/vision/tree/main/references/detection#ssdlite320-mobilenetv3-large",
"map": 21.3,
},
......
torchvision/models/efficientnet.py
View file @ f1587a20
......@@ -3,7 +3,7 @@ import math
import warnings
from dataclasses import dataclass
from functools import partial
from typing import Any, Callable, Optional, List, Sequence, Tuple, Union
from typing import Any, Callable, Dict, Optional, List, Sequence, Tuple, Union
import torch
from torch import nn, Tensor
......@@ -429,8 +429,7 @@ def _efficientnet_conf(
return inverted_residual_setting, last_channel
_COMMON_META = {
"task": "image_classification",
_COMMON_META: Dict[str, Any] = {
"categories": _IMAGENET_CATEGORIES,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#efficientnet",
}
......@@ -438,14 +437,12 @@ _COMMON_META = {
_COMMON_META_V1 = {
**_COMMON_META,
"architecture": "EfficientNet",
"min_size": (1, 1),
}
_COMMON_META_V2 = {
**_COMMON_META,
"architecture": "EfficientNetV2",
"min_size": (33, 33),
}
......@@ -459,7 +456,6 @@ class EfficientNet_B0_Weights(WeightsEnum):
meta={
**_COMMON_META_V1,
"num_params": 5288548,
"size": (224, 224),
"acc@1": 77.692,
"acc@5": 93.532,
},
......@@ -476,7 +472,6 @@ class EfficientNet_B1_Weights(WeightsEnum):
meta={
**_COMMON_META_V1,
"num_params": 7794184,
"size": (240, 240),
"acc@1": 78.642,
"acc@5": 94.186,
},
......@@ -490,7 +485,6 @@ class EfficientNet_B1_Weights(WeightsEnum):
**_COMMON_META_V1,
"num_params": 7794184,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-lr-wd-crop-tuning",
"size": (240, 240),
"acc@1": 79.838,
"acc@5": 94.934,
},
......@@ -507,7 +501,6 @@ class EfficientNet_B2_Weights(WeightsEnum):
meta={
**_COMMON_META_V1,
"num_params": 9109994,
"size": (288, 288),
"acc@1": 80.608,
"acc@5": 95.310,
},
......@@ -524,7 +517,6 @@ class EfficientNet_B3_Weights(WeightsEnum):
meta={
**_COMMON_META_V1,
"num_params": 12233232,
"size": (300, 300),
"acc@1": 82.008,
"acc@5": 96.054,
},
......@@ -541,7 +533,6 @@ class EfficientNet_B4_Weights(WeightsEnum):
meta={
**_COMMON_META_V1,
"num_params": 19341616,
"size": (380, 380),
"acc@1": 83.384,
"acc@5": 96.594,
},
......@@ -558,7 +549,6 @@ class EfficientNet_B5_Weights(WeightsEnum):
meta={
**_COMMON_META_V1,
"num_params": 30389784,
"size": (456, 456),
"acc@1": 83.444,
"acc@5": 96.628,
},
......@@ -575,7 +565,6 @@ class EfficientNet_B6_Weights(WeightsEnum):
meta={
**_COMMON_META_V1,
"num_params": 43040704,
"size": (528, 528),
"acc@1": 84.008,
"acc@5": 96.916,
},
......@@ -592,7 +581,6 @@ class EfficientNet_B7_Weights(WeightsEnum):
meta={
**_COMMON_META_V1,
"num_params": 66347960,
"size": (600, 600),
"acc@1": 84.122,
"acc@5": 96.908,
},
......@@ -612,7 +600,6 @@ class EfficientNet_V2_S_Weights(WeightsEnum):
meta={
**_COMMON_META_V2,
"num_params": 21458488,
"size": (384, 384),
"acc@1": 84.228,
"acc@5": 96.878,
},
......@@ -632,7 +619,6 @@ class EfficientNet_V2_M_Weights(WeightsEnum):
meta={
**_COMMON_META_V2,
"num_params": 54139356,
"size": (480, 480),
"acc@1": 85.112,
"acc@5": 97.156,
},
......@@ -654,7 +640,6 @@ class EfficientNet_V2_L_Weights(WeightsEnum):
meta={
**_COMMON_META_V2,
"num_params": 118515272,
"size": (480, 480),
"acc@1": 85.808,
"acc@5": 97.788,
},
......
torchvision/models/googlenet.py
View file @ f1587a20
......@@ -280,10 +280,7 @@ class GoogLeNet_Weights(WeightsEnum):
url="https://download.pytorch.org/models/googlenet-1378be20.pth",
transforms=partial(ImageClassification, crop_size=224),
meta={
"task": "image_classification",
"architecture": "GoogLeNet",
"num_params": 6624904,
"size": (224, 224),
"min_size": (15, 15),
"categories": _IMAGENET_CATEGORIES,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#googlenet",
......
torchvision/models/inception.py
View file @ f1587a20
......@@ -412,10 +412,7 @@ class Inception_V3_Weights(WeightsEnum):
url="https://download.pytorch.org/models/inception_v3_google-0cc3c7bd.pth",
transforms=partial(ImageClassification, crop_size=299, resize_size=342),
meta={
"task": "image_classification",
"architecture": "InceptionV3",
"num_params": 27161264,
"size": (299, 299),
"min_size": (75, 75),
"categories": _IMAGENET_CATEGORIES,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#inception-v3",
......
torchvision/models/mnasnet.py
View file @ f1587a20
......@@ -212,9 +212,6 @@ class MNASNet(torch.nn.Module):
_COMMON_META = {
"task": "image_classification",
"architecture": "MNASNet",
"size": (224, 224),
"min_size": (1, 1),
"categories": _IMAGENET_CATEGORIES,
"recipe": "https://github.com/1e100/mnasnet_trainer",
......
torchvision/models/mobilenetv2.py
View file @ f1587a20
......@@ -195,10 +195,7 @@ class MobileNetV2(nn.Module):
_COMMON_META = {
"task": "image_classification",
"architecture": "MobileNetV2",
"num_params": 3504872,
"size": (224, 224),
"min_size": (1, 1),
"categories": _IMAGENET_CATEGORIES,
}
......
torchvision/models/mobilenetv3.py
View file @ f1587a20
......@@ -304,9 +304,6 @@ def _mobilenet_v3(
_COMMON_META = {
"task": "image_classification",
"architecture": "MobileNetV3",
"size": (224, 224),
"min_size": (1, 1),
"categories": _IMAGENET_CATEGORIES,
}
......
torchvision/models/optical_flow/raft.py
View file @ f1587a20
......@@ -512,8 +512,7 @@ class RAFT(nn.Module):
_COMMON_META = {
"task": "optical_flow",
"architecture": "RAFT",
"min_size": (128, 128),
}
......
torchvision/models/quantization/googlenet.py
View file @ f1587a20
......@@ -111,14 +111,10 @@ class GoogLeNet_QuantizedWeights(WeightsEnum):
url="https://download.pytorch.org/models/quantized/googlenet_fbgemm-c00238cf.pth",
transforms=partial(ImageClassification, crop_size=224),
meta={
"task": "image_classification",
"architecture": "GoogLeNet",
"num_params": 6624904,
"size": (224, 224),
"min_size": (15, 15),
"categories": _IMAGENET_CATEGORIES,
"backend": "fbgemm",
"quantization": "Post Training Quantization",
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-models",
"unquantized": GoogLeNet_Weights.IMAGENET1K_V1,
"acc@1": 69.826,
......
torchvision/models/quantization/inception.py
View file @ f1587a20
......@@ -177,14 +177,10 @@ class Inception_V3_QuantizedWeights(WeightsEnum):
url="https://download.pytorch.org/models/quantized/inception_v3_google_fbgemm-71447a44.pth",
transforms=partial(ImageClassification, crop_size=299, resize_size=342),
meta={
"task": "image_classification",
"architecture": "InceptionV3",
"num_params": 27161264,
"size": (299, 299),
"min_size": (75, 75),
"categories": _IMAGENET_CATEGORIES,
"backend": "fbgemm",
"quantization": "Post Training Quantization",
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#post-training-quantized-models",
"unquantized": Inception_V3_Weights.IMAGENET1K_V1,
"acc@1": 77.176,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment