解决torchvision.models模块不存在ViT_B_16_Weights属性的报错问题

错误原因 #

你的torchvision版本过低，ViT系列模型（包括ViT_B_16_Weights权重类和vit_b_16模型构建函数）是在torchvision 0.13.0版本才正式纳入官方库的，旧版本不存在这些属性。另外原代码中vit_B_16的写法存在大小写错误，正确应为小写的vit_b_16。

解决方法 #

1. 升级torchvision到兼容版本

torch和torchvision版本需严格匹配，避免出现兼容性问题：

• 使用pip升级：

  # 升级到最新稳定版（自动匹配当前torch版本）
  pip install --upgrade torchvision
  # 或指定具体兼容版本（例如torch 2.0对应torchvision 0.15.0）
  pip install torchvision==0.15.0 --force-reinstall
  

• 使用conda升级：

  conda update torchvision
  # 或指定版本
  conda install torchvision=0.15.0
  

升级后验证版本：

import torchvision
print(torchvision.__version__)  # 输出需≥0.13.0

修正后的代码：

import torch
import torchvision

from torch import nn
from torchvision import transforms

device = "cuda" if torch.cuda.is_available() else "cpu"

pretrained_vit_weights = torchvision.models.ViT_B_16_Weights.DEFAULT 
pretrained_vit = torchvision.models.vit_b_16(weights=pretrained_vit_weights).to(device)

2. 无法升级时使用第三方实现

如果受环境限制无法升级torchvision，可以使用timm库（专门的视觉模型库）来调用ViT：

pip install timm

调用代码示例：

import timm
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
pretrained_vit = timm.create_model("vit_base_patch16_224", pretrained=True).to(device)

内容的提问来源于stack exchange，提问作者Pranav Dubal