Search code examples
pytorchtorchvisionvision-transformer

module 'torchvision.models' has no attribute 'ViT_B_16_Weights'

I am unable to import ViT_B_16 from torchvision.

I have the below code:

import torch
import torchvision

from torch import nn
from torchvision import transforms

pretrained_vit_weights = torchvision.models.ViT_B_16_Weights.DEFAULT 

pretrained_vit = torchvision.models.vit_B_16(weights=pretrained_vit_weights).to(device)

and the error is:

AttributeError                            Traceback (most recent call last)
      1 # 1. Get pretrained weights for ViT-Base
----> 2 pretrained_vit_weights = torchvision.models.ViT_B_16_Weights.DEFAULT
      3 
      4 # 2. Setup a ViT model instance with pretrained weights
      5 pretrained_vit = torchvision.models.vit_B_16(weights=pretrained_vit_weights).to(device)

AttributeError: module 'torchvision.models' has no attribute 'ViT_B_16_Weights'

resolve the error of module 'torchvision.models' has no attribute 'ViT_B_16_Weights'
Solution

  • In torch.models, the function to initialize a ViT-B-16 architecture is not vit_B_16 but vit_b_16, lowercase "b".

    As @simeonovich and @Dr. Snoopy discussed in the comments, you have on the other hand ViT_B_16_Weights (capital letters) referring to the weight enum, it refers to the different pretrained_weights: IMAGENET1K_V1 (also DEFAULT), IMAGENET1K_SWAG_E2E_V1, and IMAGENET1K_SWAG_LINEAR_V1.

    So in summary: