Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "93c64ea1b1941bc6077ff1cf48c8eac902052f77"
Unverified Commit b5aa0915 authored by Vasilis Vryniotis's avatar Vasilis Vryniotis Committed by GitHub
Browse files

Improved meta-data for models (#5170) 

* Improved meta-data for models.

* Addressing comments from code-review.

* Add parameter count.

* Fix linter.
parent 5dc61cb0
Show whitespace changes
Inline Side-by-side
Showing with 168 additions and 4 deletions
torchvision/prototype/models/quantization/mobilenetv3.py
View file @ b5aa0915
......@@ -61,6 +61,10 @@ 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,
"num_params": 5483032,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
......
torchvision/prototype/models/quantization/resnet.py
View file @ b5aa0915
......@@ -54,6 +54,7 @@ def _resnet(
_COMMON_META = {
"task": "image_classification",
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
......@@ -69,6 +70,9 @@ class ResNet18_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 11689512,
"unquantized": ResNet18_Weights.ImageNet1K_V1,
"acc@1": 69.494,
"acc@5": 88.882,
......@@ -83,6 +87,9 @@ class ResNet50_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 25557032,
"unquantized": ResNet50_Weights.ImageNet1K_V1,
"acc@1": 75.920,
"acc@5": 92.814,
......@@ -93,6 +100,9 @@ class ResNet50_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 25557032,
"unquantized": ResNet50_Weights.ImageNet1K_V2,
"acc@1": 80.282,
"acc@5": 94.976,
......@@ -107,6 +117,9 @@ class ResNeXt101_32X8D_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNeXt",
"publication_year": 2016,
"num_params": 88791336,
"unquantized": ResNeXt101_32X8D_Weights.ImageNet1K_V1,
"acc@1": 78.986,
"acc@5": 94.480,
......@@ -117,6 +130,9 @@ class ResNeXt101_32X8D_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "ResNeXt",
"publication_year": 2016,
"num_params": 88791336,
"unquantized": ResNeXt101_32X8D_Weights.ImageNet1K_V2,
"acc@1": 82.574,
"acc@5": 96.132,
......
torchvision/prototype/models/quantization/shufflenetv2.py
View file @ b5aa0915
......@@ -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,
......@@ -66,6 +69,7 @@ class ShuffleNet_V2_X0_5_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 1366792,
"unquantized": ShuffleNet_V2_X0_5_Weights.ImageNet1K_V1,
"acc@1": 57.972,
"acc@5": 79.780,
......@@ -80,6 +84,7 @@ class ShuffleNet_V2_X1_0_QuantizedWeights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 2278604,
"unquantized": ShuffleNet_V2_X1_0_Weights.ImageNet1K_V1,
"acc@1": 68.360,
"acc@5": 87.582,
......
torchvision/prototype/models/regnet.py
View file @ b5aa0915
......@@ -43,7 +43,14 @@ __all__ = [
"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(
......@@ -70,6 +77,7 @@ class RegNet_Y_400MF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 4344144,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models",
"acc@1": 74.046,
"acc@5": 91.716,
......@@ -80,6 +88,7 @@ class RegNet_Y_400MF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 4344144,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 75.804,
"acc@5": 92.742,
......@@ -94,6 +103,7 @@ class RegNet_Y_800MF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 6432512,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models",
"acc@1": 76.420,
"acc@5": 93.136,
......@@ -104,6 +114,7 @@ class RegNet_Y_800MF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 6432512,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 78.828,
"acc@5": 94.502,
......@@ -118,6 +129,7 @@ class RegNet_Y_1_6GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 11202430,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models",
"acc@1": 77.950,
"acc@5": 93.966,
......@@ -128,6 +140,7 @@ class RegNet_Y_1_6GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 11202430,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 80.876,
"acc@5": 95.444,
......@@ -142,6 +155,7 @@ class RegNet_Y_3_2GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 19436338,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models",
"acc@1": 78.948,
"acc@5": 94.576,
......@@ -152,6 +166,7 @@ class RegNet_Y_3_2GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 19436338,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 81.982,
"acc@5": 95.972,
......@@ -166,6 +181,7 @@ class RegNet_Y_8GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 39381472,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models",
"acc@1": 80.032,
"acc@5": 95.048,
......@@ -176,6 +192,7 @@ class RegNet_Y_8GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 39381472,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 82.828,
"acc@5": 96.330,
......@@ -190,6 +207,7 @@ class RegNet_Y_16GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 83590140,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#large-models",
"acc@1": 80.424,
"acc@5": 95.240,
......@@ -200,6 +218,7 @@ class RegNet_Y_16GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 83590140,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 82.886,
"acc@5": 96.328,
......@@ -214,6 +233,7 @@ class RegNet_Y_32GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 145046770,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#large-models",
"acc@1": 80.878,
"acc@5": 95.340,
......@@ -224,6 +244,7 @@ class RegNet_Y_32GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 145046770,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 83.368,
"acc@5": 96.498,
......@@ -238,6 +259,7 @@ class RegNet_X_400MF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 5495976,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models",
"acc@1": 72.834,
"acc@5": 90.950,
......@@ -248,6 +270,7 @@ class RegNet_X_400MF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 5495976,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres",
"acc@1": 74.864,
"acc@5": 92.322,
......@@ -262,6 +285,7 @@ class RegNet_X_800MF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 7259656,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models",
"acc@1": 75.212,
"acc@5": 92.348,
......@@ -272,6 +296,7 @@ class RegNet_X_800MF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 7259656,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres",
"acc@1": 77.522,
"acc@5": 93.826,
......@@ -286,6 +311,7 @@ class RegNet_X_1_6GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 9190136,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#small-models",
"acc@1": 77.040,
"acc@5": 93.440,
......@@ -296,6 +322,7 @@ class RegNet_X_1_6GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 9190136,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres",
"acc@1": 79.668,
"acc@5": 94.922,
......@@ -310,6 +337,7 @@ class RegNet_X_3_2GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 15296552,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models",
"acc@1": 78.364,
"acc@5": 93.992,
......@@ -320,6 +348,7 @@ class RegNet_X_3_2GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 15296552,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 81.196,
"acc@5": 95.430,
......@@ -334,6 +363,7 @@ class RegNet_X_8GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 39572648,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models",
"acc@1": 79.344,
"acc@5": 94.686,
......@@ -344,6 +374,7 @@ class RegNet_X_8GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 39572648,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 81.682,
"acc@5": 95.678,
......@@ -358,6 +389,7 @@ class RegNet_X_16GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 54278536,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#medium-models",
"acc@1": 80.058,
"acc@5": 94.944,
......@@ -368,6 +400,7 @@ class RegNet_X_16GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 54278536,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 82.716,
"acc@5": 96.196,
......@@ -382,6 +415,7 @@ class RegNet_X_32GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 107811560,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#large-models",
"acc@1": 80.622,
"acc@5": 95.248,
......@@ -392,6 +426,7 @@ class RegNet_X_32GF_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"num_params": 107811560,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 83.014,
"acc@5": 96.288,
......
torchvision/prototype/models/resnet.py
View file @ b5aa0915
......@@ -51,7 +51,12 @@ def _resnet(
return model
_COMMON_META = {"size": (224, 224), "categories": _IMAGENET_CATEGORIES, "interpolation": InterpolationMode.BILINEAR}
_COMMON_META = {
"task": "image_classification",
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
}
class ResNet18_Weights(WeightsEnum):
......@@ -60,6 +65,9 @@ class ResNet18_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 11689512,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
"acc@1": 69.758,
"acc@5": 89.078,
......@@ -74,6 +82,9 @@ class ResNet34_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 21797672,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
"acc@1": 73.314,
"acc@5": 91.420,
......@@ -88,6 +99,9 @@ class ResNet50_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 25557032,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
"acc@1": 76.130,
"acc@5": 92.862,
......@@ -98,6 +112,9 @@ class ResNet50_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 25557032,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres",
"acc@1": 80.674,
"acc@5": 95.166,
......@@ -112,6 +129,9 @@ class ResNet101_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 44549160,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
"acc@1": 77.374,
"acc@5": 93.546,
......@@ -122,6 +142,9 @@ class ResNet101_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 44549160,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 81.886,
"acc@5": 95.780,
......@@ -136,6 +159,9 @@ class ResNet152_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 60192808,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnet",
"acc@1": 78.312,
"acc@5": 94.046,
......@@ -146,6 +172,9 @@ class ResNet152_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "ResNet",
"publication_year": 2015,
"num_params": 60192808,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 82.284,
"acc@5": 96.002,
......@@ -160,6 +189,9 @@ class ResNeXt50_32X4D_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNeXt",
"publication_year": 2016,
"num_params": 25028904,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnext",
"acc@1": 77.618,
"acc@5": 93.698,
......@@ -170,6 +202,9 @@ class ResNeXt50_32X4D_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "ResNeXt",
"publication_year": 2016,
"num_params": 25028904,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 81.198,
"acc@5": 95.340,
......@@ -184,6 +219,9 @@ class ResNeXt101_32X8D_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "ResNeXt",
"publication_year": 2016,
"num_params": 88791336,
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#resnext",
"acc@1": 79.312,
"acc@5": 94.526,
......@@ -194,6 +232,9 @@ class ResNeXt101_32X8D_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "ResNeXt",
"publication_year": 2016,
"num_params": 88791336,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres",
"acc@1": 82.834,
"acc@5": 96.228,
......@@ -208,6 +249,9 @@ class Wide_ResNet50_2_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "WideResNet",
"publication_year": 2016,
"num_params": 68883240,
"recipe": "https://github.com/pytorch/vision/pull/912#issue-445437439",
"acc@1": 78.468,
"acc@5": 94.086,
......@@ -218,6 +262,9 @@ class Wide_ResNet50_2_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "WideResNet",
"publication_year": 2016,
"num_params": 68883240,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-fixres",
"acc@1": 81.602,
"acc@5": 95.758,
......@@ -232,6 +279,9 @@ class Wide_ResNet101_2_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"architecture": "WideResNet",
"publication_year": 2016,
"num_params": 126886696,
"recipe": "https://github.com/pytorch/vision/pull/912#issue-445437439",
"acc@1": 78.848,
"acc@5": 94.284,
......@@ -242,6 +292,9 @@ class Wide_ResNet101_2_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
meta={
**_COMMON_META,
"architecture": "WideResNet",
"publication_year": 2016,
"num_params": 126886696,
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe",
"acc@1": 82.510,
"acc@5": 96.020,
......
torchvision/prototype/models/segmentation/deeplabv3.py
View file @ b5aa0915
......@@ -25,6 +25,9 @@ __all__ = [
_COMMON_META = {
"task": "image_semantic_segmentation",
"architecture": "DeepLabV3",
"publication_year": 2017,
"categories": _VOC_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
}
......@@ -36,6 +39,7 @@ class DeepLabV3_ResNet50_Weights(WeightsEnum):
transforms=partial(VocEval, resize_size=520),
meta={
**_COMMON_META,
"num_params": 42004074,
"recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#deeplabv3_resnet50",
"mIoU": 66.4,
"acc": 92.4,
......@@ -50,6 +54,7 @@ class DeepLabV3_ResNet101_Weights(WeightsEnum):
transforms=partial(VocEval, resize_size=520),
meta={
**_COMMON_META,
"num_params": 60996202,
"recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#fcn_resnet101",
"mIoU": 67.4,
"acc": 92.4,
......@@ -64,6 +69,7 @@ class DeepLabV3_MobileNet_V3_Large_Weights(WeightsEnum):
transforms=partial(VocEval, resize_size=520),
meta={
**_COMMON_META,
"num_params": 11029328,
"recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#deeplabv3_mobilenet_v3_large",
"mIoU": 60.3,
"acc": 91.2,
......
torchvision/prototype/models/segmentation/fcn.py
View file @ b5aa0915
......@@ -15,6 +15,9 @@ __all__ = ["FCN", "FCN_ResNet50_Weights", "FCN_ResNet101_Weights", "fcn_resnet50
_COMMON_META = {
"task": "image_semantic_segmentation",
"architecture": "FCN",
"publication_year": 2014,
"categories": _VOC_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
}
......@@ -26,6 +29,7 @@ class FCN_ResNet50_Weights(WeightsEnum):
transforms=partial(VocEval, resize_size=520),
meta={
**_COMMON_META,
"num_params": 35322218,
"recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#fcn_resnet50",
"mIoU": 60.5,
"acc": 91.4,
......@@ -40,6 +44,7 @@ class FCN_ResNet101_Weights(WeightsEnum):
transforms=partial(VocEval, resize_size=520),
meta={
**_COMMON_META,
"num_params": 54314346,
"recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#deeplabv3_resnet101",
"mIoU": 63.7,
"acc": 91.9,
......
torchvision/prototype/models/segmentation/lraspp.py
View file @ b5aa0915
......@@ -19,6 +19,10 @@ class LRASPP_MobileNet_V3_Large_Weights(WeightsEnum):
url="https://download.pytorch.org/models/lraspp_mobilenet_v3_large-d234d4ea.pth",
transforms=partial(VocEval, resize_size=520),
meta={
"task": "image_semantic_segmentation",
"architecture": "LRASPP",
"publication_year": 2019,
"num_params": 3221538,
"categories": _VOC_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
"recipe": "https://github.com/pytorch/vision/tree/main/references/segmentation#lraspp_mobilenet_v3_large",
......
torchvision/prototype/models/shufflenetv2.py
View file @ b5aa0915
......@@ -41,6 +41,9 @@ def _shufflenetv2(
_COMMON_META = {
"task": "image_classification",
"architecture": "ShuffleNetV2",
"publication_year": 2018,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
......@@ -54,6 +57,7 @@ class ShuffleNet_V2_X0_5_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 1366792,
"acc@1": 69.362,
"acc@5": 88.316,
},
......@@ -67,6 +71,7 @@ class ShuffleNet_V2_X1_0_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 2278604,
"acc@1": 60.552,
"acc@5": 81.746,
},
......
torchvision/prototype/models/squeezenet.py
View file @ b5aa0915
......@@ -14,6 +14,9 @@ __all__ = ["SqueezeNet", "SqueezeNet1_0_Weights", "SqueezeNet1_1_Weights", "sque
_COMMON_META = {
"task": "image_classification",
"architecture": "SqueezeNet",
"publication_year": 2016,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
......@@ -27,6 +30,7 @@ class SqueezeNet1_0_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 1248424,
"acc@1": 58.092,
"acc@5": 80.420,
},
......@@ -40,6 +44,7 @@ class SqueezeNet1_1_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 1235496,
"acc@1": 58.178,
"acc@5": 80.624,
},
......
torchvision/prototype/models/vgg.py
View file @ b5aa0915
......@@ -41,6 +41,9 @@ def _vgg(cfg: str, batch_norm: bool, weights: Optional[WeightsEnum], progress: b
_COMMON_META = {
"task": "image_classification",
"architecture": "VGG",
"publication_year": 2014,
"size": (224, 224),
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
......@@ -54,6 +57,7 @@ class VGG11_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 132863336,
"acc@1": 69.020,
"acc@5": 88.628,
},
......@@ -67,6 +71,7 @@ class VGG11_BN_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 132868840,
"acc@1": 70.370,
"acc@5": 89.810,
},
......@@ -80,6 +85,7 @@ class VGG13_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 133047848,
"acc@1": 69.928,
"acc@5": 89.246,
},
......@@ -93,6 +99,7 @@ class VGG13_BN_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 133053736,
"acc@1": 71.586,
"acc@5": 90.374,
},
......@@ -106,6 +113,7 @@ class VGG16_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 138357544,
"acc@1": 71.592,
"acc@5": 90.382,
},
......@@ -119,9 +127,9 @@ class VGG16_Weights(WeightsEnum):
ImageNetEval, crop_size=224, mean=(0.48235, 0.45882, 0.40784), std=(1.0 / 255.0, 1.0 / 255.0, 1.0 / 255.0)
),
meta={
"size": (224, 224),
**_COMMON_META,
"num_params": 138357544,
"categories": None,
"interpolation": InterpolationMode.BILINEAR,
"recipe": "https://github.com/amdegroot/ssd.pytorch#training-ssd",
"acc@1": float("nan"),
"acc@5": float("nan"),
......@@ -136,6 +144,7 @@ class VGG16_BN_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 138365992,
"acc@1": 73.360,
"acc@5": 91.516,
},
......@@ -149,6 +158,7 @@ class VGG19_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 143667240,
"acc@1": 72.376,
"acc@5": 90.876,
},
......@@ -162,6 +172,7 @@ class VGG19_BN_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 143678248,
"acc@1": 74.218,
"acc@5": 91.842,
},
......
torchvision/prototype/models/video/resnet.py
View file @ b5aa0915
......@@ -52,6 +52,8 @@ def _video_resnet(
_COMMON_META = {
"task": "video_classification",
"publication_year": 2017,
"size": (112, 112),
"categories": _KINETICS400_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
......@@ -65,6 +67,8 @@ class R3D_18_Weights(WeightsEnum):
transforms=partial(Kinect400Eval, crop_size=(112, 112), resize_size=(128, 171)),
meta={
**_COMMON_META,
"architecture": "R3D",
"num_params": 33371472,
"acc@1": 52.75,
"acc@5": 75.45,
},
......@@ -78,6 +82,8 @@ class MC3_18_Weights(WeightsEnum):
transforms=partial(Kinect400Eval, crop_size=(112, 112), resize_size=(128, 171)),
meta={
**_COMMON_META,
"architecture": "MC3",
"num_params": 11695440,
"acc@1": 53.90,
"acc@5": 76.29,
},
......@@ -91,6 +97,8 @@ class R2Plus1D_18_Weights(WeightsEnum):
transforms=partial(Kinect400Eval, crop_size=(112, 112), resize_size=(128, 171)),
meta={
**_COMMON_META,
"architecture": "R(2+1)D",
"num_params": 31505325,
"acc@1": 57.50,
"acc@5": 78.81,
},
......
torchvision/prototype/models/vision_transformer.py
View file @ b5aa0915
......@@ -236,6 +236,9 @@ class VisionTransformer(nn.Module):
_COMMON_META = {
"task": "image_classification",
"architecture": "ViT",
"publication_year": 2020,
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BILINEAR,
}
......@@ -247,6 +250,7 @@ class ViT_B_16_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 86567656,
"size": (224, 224),
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#vit_b_16",
"acc@1": 81.072,
......@@ -262,6 +266,7 @@ class ViT_B_32_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 88224232,
"size": (224, 224),
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#vit_b_32",
"acc@1": 75.912,
......@@ -277,6 +282,7 @@ class ViT_L_16_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224, resize_size=242),
meta={
**_COMMON_META,
"num_params": 304326632,
"size": (224, 224),
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#vit_l_16",
"acc@1": 79.662,
......@@ -292,6 +298,7 @@ class ViT_L_32_Weights(WeightsEnum):
transforms=partial(ImageNetEval, crop_size=224),
meta={
**_COMMON_META,
"num_params": 306535400,
"size": (224, 224),
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#vit_l_32",
"acc@1": 76.972,
......
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