在 Hugging Face 上使用 mlx-image

mlx-image 是 Riccardo Musmeci 基於 Apple MLX 開發的影像模型庫。它嘗試複製出色的 timm，但適用於 MLX 模型。

在 Hub 上探索 mlx-image

您可以透過使用 `mlx-image` 庫名稱進行過濾來查詢 `mlx-image` 模型，如此查詢所示。還有一個開放的 mlx-vision 社群，供貢獻者轉換和釋出 MLX 格式的權重。

安裝

pip install mlx-image

模型

模型權重可在 HuggingFace 上的 mlx-vision 社群中獲取。

載入預訓練模型權重

from mlxim.model import create_model

# loading weights from HuggingFace (https://huggingface.co/mlx-vision/resnet18-mlxim)
model = create_model("resnet18") # pretrained weights loaded from HF

# loading weights from local file
model = create_model("resnet18", weights="path/to/resnet18/model.safetensors")

列出所有可用模型

from mlxim.model import list_models
list_models()

ImageNet-1K 結果

請訪問 results-imagenet-1k.csv，檢視每個轉換為 mlx-image 的模型及其在不同設定下在 ImageNet-1K 上的效能。

簡而言之，效能與 PyTorch 實現中的原始模型相當。

與 PyTorch 和其他熟悉工具的相似性

mlx-image 嘗試儘可能接近 PyTorch。

DataLoader -> 您可以定義自己的 collate_fn，也可以使用 num_workers 來加快資料載入速度。
Dataset -> mlx-image 已支援 LabelFolderDataset（與舊的 PyTorch ImageFolder 類似）和 FolderDataset（一個包含影像的通用資料夾）。
ModelCheckpoint -> 跟蹤最佳模型並將其儲存到磁碟（類似於 PyTorchLightning）。它還建議提前停止。

訓練

訓練類似於 PyTorch。這是一個如何訓練模型的示例：

import mlx.nn as nn
import mlx.optimizers as optim
from mlxim.model import create_model
from mlxim.data import LabelFolderDataset, DataLoader

train_dataset = LabelFolderDataset(
    root_dir="path/to/train",
    class_map={0: "class_0", 1: "class_1", 2: ["class_2", "class_3"]}
)
train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=4
)
model = create_model("resnet18") # pretrained weights loaded from HF
optimizer = optim.Adam(learning_rate=1e-3)

def train_step(model, inputs, targets):
    logits = model(inputs)
    loss = mx.mean(nn.losses.cross_entropy(logits, target))
    return loss

model.train()
for epoch in range(10):
    for batch in train_loader:
        x, target = batch
        train_step_fn = nn.value_and_grad(model, train_step)
        loss, grads = train_step_fn(x, target)
        optimizer.update(model, grads)
        mx.eval(model.state, optimizer.state)

附加資源

聯絡方式

如果您有任何問題，請傳送電子郵件至 riccardomusmeci92@gmail.com。

< > 在 GitHub 上更新

中心