利用vit模型提取图片的特征向量代码

以下是使用vit模型提取图片特征向量的代码：

import torch
import torchvision.transforms as transforms
import torch.nn.functional as F
from PIL import Image
from torchvision.models import vision_transformer

# 加载预训练的vit模型
model = vision_transformer(pretrained=True)

# 图片预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载图片
img = Image.open('example.jpg')

# 预处理图片
img = transform(img)

# 添加batch维度
img = img.unsqueeze(0)

# 提取特征向量
with torch.no_grad():
    features = model(img)

# 对特征向量进行归一化
features = F.normalize(features, p=2, dim=1)

# 输出特征向量
print(features)

希望这个代码对你有所帮助！