许可证：apache-2.0
标签：

• mlx
• mlx-image
• 视觉
• 图像分类
  数据集：
• imagenet-1k
  库名称：mlx-image

vit_base_patch16_224.dino

基于Vision Transformer架构的图像分类模型，使用DINO方法在ImageNet-1k数据集上训练而成。

该模型以自监督方式在ImageNet-1k数据集上训练，仅训练了骨干网络，未训练分类头部。

免责声明：此模型为将PyTorch模型权重移植至Apple MLX框架的版本。

使用方法

pip install mlx-image  

以下为图像分类的调用示例：

from mlxim.model import create_model  
from mlxim.io import read_rgb  
from mlxim.transform import ImageNetTransform  

transform = ImageNetTransform(train=False, img_size=224)  
x = transform(read_rgb("cat.png"))  
x = mx.expand_dims(x, 0)  

model = create_model("vit_base_patch16_224.dino")  
model.eval()  

logits, attn_masks = model(x, attn_masks=True)  

也可提取分类头之前的层特征：

# 第一种方式  
model = create_model("vit_base_patch16_224.dino", num_classes=0)  
model.eval()  
embeds = model(x)  

# 第二种方式  
model = create_model("vit_base_patch16_224.dino")  
model.eval()  
embeds, attn_masks = model.get_features(x)  

注意力热图

通过模型返回的attn_masks可可视化注意力热图，详见mlx-image的示例笔记本。