Usando mlx-image no Hugging Face

mlx-image é uma biblioteca de modelos de imagem desenvolvida por Riccardo Musmeci construída sobre o MLX da Apple. Ela tenta replicar o excelente timm, mas para modelos MLX.

Explorando o mlx-image no Hub

Você pode encontrar modelos de mlx-image filtrando pelo nome da biblioteca mlx-image, como em esta consulta. Também existe uma comunidade aberta de mlx-vision para contribuidores que convertem e publicam pesos no formato MLX.

Instalação

pip install mlx-image

Modelos

Os pesos do modelo estão disponíveis na comunidade mlx-vision no Hugging Face.

Para carregar um modelo com pesos pré-treinados:

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

Para listar todos os modelos disponíveis:

from mlxim.model import list_models
list_models()

Resultados do ImageNet-1K

Acesse results-imagenet-1k.csv para verificar todos os modelos convertidos para mlx-image e seu desempenho no ImageNet-1K com diferentes configurações.

O desempenho do TL;DR é comparável ao dos modelos originais das implementações do PyTorch.

Similaridade com o PyTorch e outras ferramentas familiares

mlx-image tenta ser o mais próximo possível do PyTorch:

• DataLoader -> você pode definir sua própria collate_fn e também usar num_workers para acelerar o carregamento de dados

• Dataset -> mlx-image já suporta LabelFolderDataset (o bom e velho ImageFolder do PyTorch) e FolderDataset (uma pasta genérica com imagens)

• ModelCheckpoint -> acompanha o melhor modelo e o salva em disco (similar ao PyTorchLightning). Também sugere o encerramento antecipado

Treinamento

O treinamento é semelhante ao do PyTorch. Aqui está um exemplo de como treinar um modelo:

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)

Recursos adicionais

• repositório mlx-image
• comunidade mlx-vision

Contato

Se você tiver alguma dúvida, por favor, envie um e-mail para riccardomusmeci92@gmail.com.