Utiliser mlx-image sur Hugging Face

mlx-image est une bibliothèque de modèles d’images développée par Riccardo Musmeci et basée sur MLX d’Apple. Elle tente de reproduire les qualités de timm, mais pour les modèles MLX.

Explorer mlx-image sur le Hub

Vous pouvez trouver des modèles mlx-image en filtrant par le nom de la bibliothèque mlx-image, comme dans cette requête. Il existe également une communauté ouverte mlx-vision pour les contributeurs qui convertissent et publient des poids au format MLX.

Installation

pip install mlx-image

Modèles

Les poids du modèle sont disponibles sur la communauté mlx-vision de Hugging Face.

Pour charger un modèle avec des poids pré-entraînés :

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")

Pour lister tous les modèles disponibles :

from mlxim.model import list_models
list_models()

Résultats sur ImageNet-1K

Accédez à results-imagenet-1k.csv pour consulter chaque modèle converti en mlx-image et ses performances sur ImageNet-1K avec différentes configurations.

TL;DR les performances sont comparables à celles des modèles originaux issus des implémentations PyTorch.

Similarité avec PyTorch et d'autres outils familiers

mlx-image essaie d’être aussi proche que possible de PyTorch :

• DataLoader -> vous pouvez définir votre propre collate_fn et utiliser num_workers pour accélérer le chargement des données

• Dataset -> mlx-image prend déjà en charge LabelFolderDataset (le bon vieux ImageFolder de PyTorch) et FolderDataset (un dossier générique contenant des images)

• ModelCheckpoint -> suit le meilleur modèle et l'enregistre sur le disque (similaire à PyTorchLightning). Il suggère également l'arrêt anticipé

Entraînement

L'entraînement est similaire à PyTorch. Voici un exemple de la façon d'entraîner un modèle :

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)

Ressources supplémentaires

• dépôt mlx-image
• communauté mlx-vision

Contact

Si vous avez des questions, veuillez envoyer un e-mail à riccardomusmeci92@gmail.com.