Skip to content

Commit 2b5ab1b

Browse files
authored
Adding improved MobileNetV2 weights (#5560)
* Adding improved MobileNetV2 weights * Fix eval sizes. * Update accuracy with batch-size 1 run
1 parent 358fc6d commit 2b5ab1b

File tree

1 file changed

+24
-9
lines changed

1 file changed

+24
-9
lines changed

torchvision/prototype/models/mobilenetv2.py

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,25 +13,40 @@
1313
__all__ = ["MobileNetV2", "MobileNet_V2_Weights", "mobilenet_v2"]
1414

1515

16+
_COMMON_META = {
17+
"task": "image_classification",
18+
"architecture": "MobileNetV2",
19+
"publication_year": 2018,
20+
"num_params": 3504872,
21+
"size": (224, 224),
22+
"min_size": (1, 1),
23+
"categories": _IMAGENET_CATEGORIES,
24+
"interpolation": InterpolationMode.BILINEAR,
25+
}
26+
27+
1628
class MobileNet_V2_Weights(WeightsEnum):
1729
IMAGENET1K_V1 = Weights(
1830
url="https://download.pytorch.org/models/mobilenet_v2-b0353104.pth",
1931
transforms=partial(ImageNetEval, crop_size=224),
2032
meta={
21-
"task": "image_classification",
22-
"architecture": "MobileNetV2",
23-
"publication_year": 2018,
24-
"num_params": 3504872,
25-
"size": (224, 224),
26-
"min_size": (1, 1),
27-
"categories": _IMAGENET_CATEGORIES,
28-
"interpolation": InterpolationMode.BILINEAR,
33+
**_COMMON_META,
2934
"recipe": "https://github.com/pytorch/vision/tree/main/references/classification#mobilenetv2",
3035
"acc@1": 71.878,
3136
"acc@5": 90.286,
3237
},
3338
)
34-
DEFAULT = IMAGENET1K_V1
39+
IMAGENET1K_V2 = Weights(
40+
url="https://download.pytorch.org/models/mobilenet_v2-7ebf99e0.pth",
41+
transforms=partial(ImageNetEval, crop_size=224, resize_size=232),
42+
meta={
43+
**_COMMON_META,
44+
"recipe": "https://github.com/pytorch/vision/issues/3995#new-recipe-with-reg-tuning",
45+
"acc@1": 72.154,
46+
"acc@5": 90.822,
47+
},
48+
)
49+
DEFAULT = IMAGENET1K_V2
3550

3651

3752
@handle_legacy_interface(weights=("pretrained", MobileNet_V2_Weights.IMAGENET1K_V1))

0 commit comments

Comments
 (0)