Использование mlx-image на Hugging Face

mlx-image — это библиотека моделей изображений, разработанная Riccardo Musmeci на базе Apple MLX. Она стремится воспроизвести возможности популярной библиотеки timm, но для моделей MLX.

Исследование mlx-image на Hub

Модели mlx-image можно найти, отфильтровав по имени библиотеки mlx-image, как в этом запросе. Также существует открытое сообщество mlx-vision для участников, конвертирующих и публикующих веса в формате MLX.

Установка

pip install mlx-image

Модели

Веса модели доступны в сообществе mlx-vision на Hugging Face.

Для загрузки модели с предварительно обученными весами:

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 (хорошую и старую ImageFolder из PyTorch) и 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)

Дополнительные ресурсы

• репозиторий mlx-image
• сообщество mlx-vision

Контакты

Если у вас есть вопросы, напишите на riccardomusmeci92@gmail.com.