### 在 PyTorch 中加载并使用预训练的 ViT 模型 为了在 PyTorch 中加载并使用预训练的 Vision Transformer (ViT) 模型，可以借助 Hugging Face 的 `transformers` 库。该库提供了方便的方法来下载和初始化基于不同数据集预训练的 ViT 模型[^1]。 以下是具体的实现方法： #### 1. 安装必要的依赖项 首先需要安装 Hugging Face 的 `transformers` 和 `datasets` 库，以及 PyTorch 自身的支持包。 ```bash pip install transformers torch datasets ``` #### 2. 导入所需的类和函数 导入 `ViTFeatureExtractor` 和 `ViTForImageClassification` 类分别用于处理输入图像和加载预训练模型。 ```python from PIL import Image import requests from transformers import ViTFeatureExtractor, ViTForImageClassification import torch ``` #### 3. 下载或指定本地路径加载预训练模型 如果要从远程仓库获取预训练模型，则可以直接通过名称调用；如果是本地存储的配置文件夹，则需提供确切路径[^1]。 ##### 方法一：在线加载官方预训练权重 ```python model_name_or_path = 'google/vit-base-patch16-224' # 可替换为其他可用版本 feature_extractor = ViTFeatureExtractor.from_pretrained(model_name_or_path) model = ViTForImageClassification.from_pretrained(model_name_or_path) ``` ##### 方法二：离线模式下从自定义目录读取 当已经下载好特定的预训练参数至某位置时，可按照下面方式完成实例化： ```python local_model_dir = '/path/to/local/model/directory' feature_extractor = ViTFeatureExtractor.from_pretrained(local_model_dir) model = ViTForImageClassification.from_pretrained(local_model_dir) ``` 注意这里的 `/path/to/local/model/directory` 应替换成实际存放模型文件的地方，例如前次例子中的桌面路径 `/home/yejiangchen/Desktop/Codes/ViT/config/`[^1]。 #### 4. 准备输入图片 对于单幅测试图象的情况，先将其转换成适配 ViTs 输入尺寸的形式再送入网络预测类别标签。 ```python url = "http://images.cocodataset.org/val2017/000000039769.jpg" image = Image.open(requests.get(url, stream=True).raw) inputs = feature_extractor(images=image, return_tensors="pt") outputs = model(**inputs) logits = outputs.logits predicted_class_idx = logits.argmax(-1).item() print(f"Predicted class: {model.config.id2label[predicted_class_idx]}") ``` 此脚本片段展示了如何抓取一张公开 URL 图片作为样例，并经过特征提取器编码后传递给分类头得出最终推断结果[^2]。 --- ###