Sử dụng mlx-image trên Hugging Face

mlx-image là thư viện mô hình hình ảnh do Riccardo Musmeci phát triển, được xây dựng trên MLX của Apple. Thư viện này cố gắng tái hiện timm tuyệt vời, nhưng dành cho các mô hình MLX.

Khám phá mlx-image trên Hub

Bạn có thể tìm các model mlx-image bằng cách lọc theo tên thư viện mlx-image, như trong truy vấn này. Cũng có một cộng đồng mlx-vision mở cho các nhà đóng góp chuyển đổi và phát hành weights cho định dạng MLX.

Cài đặt

pip install mlx-image

Mô hình

Trọng số mô hình có sẵn trên cộng đồng mlx-vision trên Hugging Face.

Để tải một mô hình với trọng số được huấn luyện trước:

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

Để liệt kê tất cả các model khả dụng:

from mlxim.model import list_models
list_models()

Kết quả ImageNet-1K

Chuyển đến results-imagenet-1k.csv để kiểm tra mọi model đã được chuyển đổi sang mlx-image và hiệu suất của chúng trên ImageNet-1K với các cài đặt khác nhau.

Tóm tắt hiệu năng tương đương với các mô hình gốc từ các triển khai PyTorch.

Sự tương đồng với PyTorch và các công cụ quen thuộc khác

mlx-image cố gắng bám sát PyTorch nhất có thể:

• DataLoader -> bạn có thể tự định nghĩa collate_fn và cũng có thể sử dụng num_workers để tăng tốc độ tải dữ liệu

• Dataset -> mlx-image đã hỗ trợ LabelFolderDataset (chính là ImageFolder quen thuộc của PyTorch) và FolderDataset (một thư mục chung chứa các hình ảnh)

• ModelCheckpoint -> theo dõi mô hình tốt nhất và lưu nó xuống đĩa (tương tự PyTorchLightning). Nó cũng đề xuất dừng sớm

Huấn luyện

Quá trình huấn luyện tương tự như PyTorch. Dưới đây là một ví dụ về cách huấn luyện một mô hình:

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)

Tài nguyên bổ sung

• kho mlx-image
• cộng đồng mlx-vision

Liên hệ

Nếu bạn có bất kỳ câu hỏi nào, vui lòng gửi email đến riccardomusmeci92@gmail.com.