Skip to content

Add SWAG Vision Transformer Weight #5714

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 13 commits into from
Apr 1, 2022
Merged

Add SWAG Vision Transformer Weight #5714

Changes from 1 commit
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
c801bf0
Add vit_b_16_swag
YosuaMichael Mar 31, 2022
9e13f79
Better handling idiom for image_size, edit test_extended_model to han…
YosuaMichael Mar 31, 2022
1707171
Update the accuracy to the experiment result on torchvision model
YosuaMichael Mar 31, 2022
bd8b1a8
Fix typo missing underscore
YosuaMichael Mar 31, 2022
6c765a5
raise exception instead of torch._assert, add back publication year (…
YosuaMichael Mar 31, 2022
3326e88
Merge branch 'main' into add-swag-weight
YosuaMichael Apr 1, 2022
e444c5a
Add license information on meta and readme
YosuaMichael Apr 1, 2022
a6ee605
Merge branch 'add-swag-weight' of github.com:pytorch/vision into add-…
YosuaMichael Apr 1, 2022
54aa8cf
Improve wording and fix typo for pretrained model license in readme
YosuaMichael Apr 1, 2022
f9c32eb
Add vit_l_16 weight
YosuaMichael Apr 1, 2022
4cf4eff
Update README.rst
YosuaMichael Apr 1, 2022
9230f40
Update the accuracy meta on vit_l_16_swag model to result from our ex…
YosuaMichael Apr 1, 2022
c36a3ec
Merge branch 'main' into add-swag-weight
YosuaMichael Apr 1, 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
34 changes: 33 additions & 1 deletion torchvision/models/vision_transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import collections.abc as abc
import math
from collections import OrderedDict
from functools import partial
Expand Down Expand Up @@ -284,7 +285,15 @@ def _vision_transformer(
progress: bool,
**kwargs: Any,
) -> VisionTransformer:
image_size = kwargs.pop("image_size", 224)

image_size = None
if "image_size" in kwargs:
image_size = kwargs.pop("image_size", None)
if image_size is None and weights is not None and "size" in weights.meta:
image_size = weights.meta["size"]
if isinstance(image_size, abc.Sequence):
image_size = image_size[0]
image_size = image_size or 224

if weights is not None:
_ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"]))
Expand Down Expand Up @@ -313,6 +322,15 @@ def _vision_transformer(
"interpolation": InterpolationMode.BILINEAR,
}

_COMMON_SWAG_META = {
"task": "image_classification",
"architecture": "ViT",
"publication_year": 2022,
"recipe": "https://github.com/facebookresearch/SWAG",
"categories": _IMAGENET_CATEGORIES,
"interpolation": InterpolationMode.BICUBIC,
}


class ViT_B_16_Weights(WeightsEnum):
IMAGENET1K_V1 = Weights(
Expand All @@ -328,6 +346,20 @@ class ViT_B_16_Weights(WeightsEnum):
"acc@5": 95.318,
},
)
IMAGENET1K_SWAG_V1 = Weights(
url="https://download.pytorch.org/models/vit_b_16_swag-9ac1b537.pth",
transforms=partial(
ImageClassification, resize_size=384, interpolation=InterpolationMode.BICUBIC, crop_size=384
),
meta={
**_COMMON_SWAG_META,
"num_params": 86859496,
"size": (384, 384),
"min_size": (384, 384),
"acc@1": 85.29,
"acc@5": 97.65,
},
)
DEFAULT = IMAGENET1K_V1


Expand Down