Hub documentation

Hugging Faceでのmlx-imageの使用

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Hugging Faceでのmlx-imageの使用

mlx-image は、Apple MLX を基盤として Riccardo Musmeci によって開発された画像モデル用ライブラリです。timm の優れた機能を MLX モデル向けに再現することを目的としています。

Hub での mlx-image の探索

mlx-imageのモデルは、このクエリのように、mlx-imageというライブラリ名でフィルタリングすることで見つけることができます。 また、MLX形式への重みの変換と公開に取り組む貢献者のためのオープンコミュニティ mlx-vision もあります。

インストール

pip install mlx-image

モデル

モデルの重みは、Hugging Faceの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の結果

mlx-image に変換されたすべてのモデルと、異なる設定での ImageNet-1K でのパフォーマンスを確認するには、results-imagenet-1k.csv を参照してください。

TL;DR 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 までメールをお送りください。

Update on GitHub