Vit工程化应用（timm 库）

pip install timm

importtimmimporttorch from PILimportImageimportrequests from ioimportBytesIO#1.加载模型(ViTBase版本，16x16图块，在ImageNet-1k上预训练)# 设置 pretrained=True自动下载权重 model=timm.create_model('vit_base_patch16_224',pretrained=True)model.eval()# 切换到推理模式 #2.获取数据预处理流程(模型自带，确保尺寸和归一化参数正确)data_config=timm.data.resolve_model_data_config(model)transforms=timm.data.create_transform(**data_config,is_training=False)#3.读取图片(示例：从URL读取，你也可以改成本地路径)# 如果是本地文件：img=Image.open('your_image.jpg').convert('RGB')url='http://images.cocodataset.org/val2017/000000039769.jpg'# 两张猫的示例图 img=Image.open(requests.get(url,stream=True).raw).convert('RGB')#4.预处理并推理 tensor_img=transforms(img).unsqueeze(0)# 增加 batch 维度withtorch.no_grad():outputs=model(tensor_img)#5.解析结果(获取前5个预测类别)probabilities=torch.nn.functional.softmax(outputs[0],dim=0)# 下载ImageNet标签 labels_url="https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"labels=requests.get(labels_url).text.splitlines()# 打印Top5top5_prob,top5_catid=torch.topk(probabilities,5)foriinrange(top5_prob.size(0)):print(f"{labels[top5_catid[i]]}: {top5_prob[i].item():.4f}")