How to access latest torchvision.models (e.g. ViT)?

user_514149 0
Viewed 1166

I have seen in the official torchvision docs that recently vision transformers and the ConvNeXt model families have been added to the PyTorch model zoo. However, even after upgrading to latest torchvision version 0.11.3 (via pip) these new models are not available:

>>> import torchvision; torchvision.__version__
'0.11.3+cu102'
>>> import torchvision.models as models
>>> model = models.resnext50_32x4d()  # previous models work fine
>>> model = models.vit_b_16()  # vision transformers don't work
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: module 'torchvision.models' has no attribute 'vit_b_16'

Any ideas how I can access these latest model additions in PyTorch?

pytorch python
2 Answers

There are other ways of getting pytorch models besides torchvision. You should look at torch.hub for getting models from specific github repos that enabled sharing via this interface.

Additionally, you have timm -- a repository for many pytorch vision models.

for example:

# list all ViT models
timm.list_models('vit_*')
# list all convNext models
timm.list_models('convnext*')

# load ViT-B/16
vit_b_16 = timm.create_model('vit_base_patch16_224', pretrained=True)
# load conv next
convnext = timm.create_model('convnext_base', pretrained=True)
user_17144100

Event though @Shai's answer is a nice addition, my original question was how I could access the official ViT and ConvNeXt models in torchvision.models. As it turned out the answer was simply to wait. So for the records: After upgrading to latest torchvision pip package in version 0.12 I got these new models as well.

user_5141490
Related