Skip to content

Improved meta-data for models #5170

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Jan 6, 2022
Merged

Improved meta-data for models #5170

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
03141c9
Improved meta-data for models.
datumbox Jan 6, 2022
2deb7c4
Addressing comments from code-review.
datumbox Jan 6, 2022
90fd5fd
Add parameter count.
datumbox Jan 6, 2022
93d9fc4
Fix linter.
datumbox Jan 6, 2022
8cae313
Merge branch 'main' into models/enhanced_meta
datumbox Jan 6, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion test/test_prototype_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ def test_naming_conventions(model_fn):
def test_schema_meta_validation(model_fn):
classification_fields = ["size", "categories", "acc@1", "acc@5"]
defaults = {
"all": ["interpolation", "recipe"],
"all": ["task", "architecture", "publication_year", "interpolation", "recipe"],
"models": classification_fields,
"detection": ["categories", "map"],
"quantization": classification_fields + ["backend", "quantization", "unquantized"],
Expand Down
4 changes: 3 additions & 1 deletion torchvision/models/segmentation/fcn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,9 @@

class FCN(_SimpleSegmentationModel):
"""
Implements a Fully-Convolutional Network for semantic segmentation.
Implements FCN model from
`"Fully Convolutional Networks for Semantic Segmentation"
<https://arxiv.org/abs/1411.4038>`_.

Args:
backbone (nn.Module): the network used to compute the features for the model.
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/alexnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ class AlexNet_Weights(WeightsEnum):
url="https://download.pytorch.org/models/alexnet-owt-7be5be79.pth",
transforms=partial(ImageNetEval, crop_size=224),
meta={
"task": "image_classification",
"architecture": "AlexNet",
"publication_year": 2012,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/densenet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,9 @@ def _densenet(


_COMMON_META = {
"task": "image_classification",
"architecture": "DenseNet",
"publication_year": 2016,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/detection/faster_rcnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@


_COMMON_META = {
"task": "image_object_detection",
"architecture": "FasterRCNN",
"publication_year": 2015,
"categories": _COCO_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
}
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/detection/keypoint_rcnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,9 @@


_COMMON_META = {
"task": "image_object_detection",
Copy link
Contributor

@oke-aditya oke-aditya Jan 9, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

image keypoint detection. I guess?

"architecture": "FasterRCNN",
"publication_year": 2015,
"categories": _COCO_PERSON_CATEGORIES,
"keypoint_names": _COCO_PERSON_KEYPOINT_NAMES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/detection/mask_rcnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@ class MaskRCNN_ResNet50_FPN_Weights(WeightsEnum):
url="https://download.pytorch.org/models/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth",
transforms=CocoEval,
meta={
"task": "image_object_detection",
Copy link
Member

@fmassa fmassa Jan 6, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm unclear on what to put here.
Indeed, this model could also be categorized as image_instance_segmentation. Maybe having task be a list would be of help here?

Ensuring that this is adapted with paperswithcode categorization would be good as well https://paperswithcode.com/paper/mask-r-cnn

Copy link
Contributor Author

@datumbox datumbox Jan 6, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah same here. I'll leave as is for now and review in more detail on the future.

"architecture": "FasterRCNN",
"publication_year": 2015,
"categories": _COCO_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
"recipe": "https://github.com/pytorch/vision/tree/main/references/detection#mask-r-cnn",
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/detection/retinanet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,9 @@ class RetinaNet_ResNet50_FPN_Weights(WeightsEnum):
url="https://download.pytorch.org/models/retinanet_resnet50_fpn_coco-eeacb38b.pth",
transforms=CocoEval,
meta={
"task": "image_object_detection",
"architecture": "RetinaNet",
"publication_year": 2017,
"categories": _COCO_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
"recipe": "https://github.com/pytorch/vision/tree/main/references/detection#retinanet",
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/detection/ssd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,9 @@ class SSD300_VGG16_Weights(WeightsEnum):
url="https://download.pytorch.org/models/ssd300_vgg16_coco-b556d3b4.pth",
transforms=CocoEval,
meta={
"task": "image_object_detection",
"architecture": "SSD",
"publication_year": 2015,
"size": (300, 300),
"categories": _COCO_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/detection/ssdlite.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,9 @@ class SSDLite320_MobileNet_V3_Large_Weights(WeightsEnum):
url="https://download.pytorch.org/models/ssdlite320_mobilenet_v3_large_coco-a79551df.pth",
transforms=CocoEval,
meta={
"task": "image_object_detection",
"architecture": "SSDLite",
"publication_year": 2018,
"size": (320, 320),
"categories": _COCO_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/efficientnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,9 @@ def _efficientnet(


_COMMON_META = {
"task": "image_classification",
"architecture": "EfficientNet",
"publication_year": 2019,
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BICUBIC,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#efficientnet",
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/googlenet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ class GoogLeNet_Weights(WeightsEnum):
url="https://download.pytorch.org/models/googlenet-1378be20.pth",
transforms=partial(ImageNetEval, crop_size=224),
meta={
"task": "image_classification",
"architecture": "GoogLeNet",
"publication_year": 2014,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/inception.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ class Inception_V3_Weights(WeightsEnum):
url="https://download.pytorch.org/models/inception_v3_google-0cc3c7bd.pth",
transforms=partial(ImageNetEval, crop_size=299, resize_size=342),
meta={
"task": "image_classification",
"architecture": "InceptionV3",
"publication_year": 2015,
"size": (299, 299),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/mnasnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,9 @@


_COMMON_META = {
"task": "image_classification",
"architecture": "MNASNet",
"publication_year": 2018,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/mobilenetv2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ class MobileNet_V2_Weights(WeightsEnum):
url="https://download.pytorch.org/models/mobilenet_v2-b0353104.pth",
transforms=partial(ImageNetEval, crop_size=224),
meta={
"task": "image_classification",
"architecture": "MobileNetV2",
"publication_year": 2018,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/mobilenetv3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@ def _mobilenet_v3(


_COMMON_META = {
"task": "image_classification",
"architecture": "MobileNetV3",
"publication_year": 2019,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
7 changes: 6 additions & 1 deletion torchvision/prototype/models/optical_flow/raft.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,12 @@
)


_COMMON_META = {"interpolation": InterpolationMode.BILINEAR}
_COMMON_META = {
"task": "optical_flow",
"architecture": "RAFT",
"publication_year": 2020,
"interpolation": InterpolationMode.BILINEAR
}


class Raft_Large_Weights(WeightsEnum):
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/quantization/googlenet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@ class GoogLeNet_QuantizedWeights(WeightsEnum):
url="https://download.pytorch.org/models/quantized/googlenet_fbgemm-c00238cf.pth",
transforms=partial(ImageNetEval, crop_size=224),
meta={
"task": "image_classification",
"architecture": "GoogLeNet",
"publication_year": 2014,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/quantization/inception.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,9 @@ class Inception_V3_QuantizedWeights(WeightsEnum):
url="https://download.pytorch.org/models/quantized/inception_v3_google_fbgemm-71447a44.pth",
transforms=partial(ImageNetEval, crop_size=299, resize_size=342),
meta={
"task": "image_classification",
"architecture": "InceptionV3",
"publication_year": 2015,
"size": (299, 299),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/quantization/mobilenetv2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,9 @@ class MobileNet_V2_QuantizedWeights(WeightsEnum):
url="https://download.pytorch.org/models/quantized/mobilenet_v2_qnnpack_37f702c5.pth",
transforms=partial(ImageNetEval, crop_size=224),
meta={
"task": "image_classification",
"architecture": "MobileNetV2",
"publication_year": 2018,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/quantization/mobilenetv3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,9 @@ class MobileNet_V3_Large_QuantizedWeights(WeightsEnum):
url="https://download.pytorch.org/models/quantized/mobilenet_v3_large_qnnpack-5bcacf28.pth",
transforms=partial(ImageNetEval, crop_size=224),
meta={
"task": "image_classification",
"architecture": "MobileNetV3",
"publication_year": 2019,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
11 changes: 11 additions & 0 deletions torchvision/prototype/models/quantization/resnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ def _resnet(


_COMMON_META = {
"task": "image_classification",
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand All @@ -69,6 +70,8 @@ class ResNet18_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"unquantized": ResNet18_Weights.ImageNet1K_V1,
"acc@1": 69.494,
"acc@5": 88.882,
Expand All @@ -83,6 +86,8 @@ class ResNet50_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"unquantized": ResNet50_Weights.ImageNet1K_V1,
"acc@1": 75.920,
"acc@5": 92.814,
Expand All @@ -93,6 +98,8 @@ class ResNet50_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"unquantized": ResNet50_Weights.ImageNet1K_V2,
"acc@1": 80.282,
"acc@5": 94.976,
Expand All @@ -107,6 +114,8 @@ class ResNeXt101_32X8D_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNeXt",
"publication_year": 2016,
"unquantized": ResNeXt101_32X8D_Weights.ImageNet1K_V1,
"acc@1": 78.986,
"acc@5": 94.480,
Expand All @@ -117,6 +126,8 @@ class ResNeXt101_32X8D_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "ResNeXt",
"publication_year": 2016,
"unquantized": ResNeXt101_32X8D_Weights.ImageNet1K_V2,
"acc@1": 82.574,
"acc@5": 96.132,
Expand Down
3 changes: 3 additions & 0 deletions torchvision/prototype/models/quantization/shufflenetv2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,9 @@ def _shufflenetv2(


_COMMON_META = {
"task": "image_classification",
"architecture": "ShuffleNetV2",
"publication_year": 2018,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
Expand Down
9 changes: 8 additions & 1 deletion torchvision/prototype/models/regnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,14 @@
"regnet_x_32gf",
]

_COMMON_META = {"size": (224, 224), "categories": _IMAGENET_CATEGORIES, "interpolation": InterpolationMode.BILINEAR}
_COMMON_META = {
"task": "image_classification",
"architecture": "RegNet",
"publication_year": 2020,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR
}


def _regnet(
Expand Down
Loading