Vit Base Patch16 224.dino Mlxim

🚀 vit_base_patch16_224.dino

這是一個基於Vision Transformer架構的圖像分類模型，使用DINO方法在ImageNet - 1k數據集上進行訓練。

該模型在ImageNet - 1k數據集上以自監督的方式進行訓練，僅訓練了骨幹網絡，未訓練分類頭。

聲明：這是將torch模型權重移植到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)

高級用法

你也可以使用模型頭部前一層的嵌入向量：

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

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

# 第一種方式
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的notebook。

📄 許可證

本項目採用Apache - 2.0許可證。