Hub documentation
Hugging Faceでのmlx-imageの使用
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 までメールをお送りください。