Hub documentation
在 Hugging Face 上使用 mlx-image
在 Hugging Face 上使用 mlx-image
mlx-image 是由 Riccardo Musmeci 基于 Apple MLX 开发的图像模型库。它旨在复刻出色的 timm，但专为 MLX 模型设计。
在 Hub 上探索 mlx-image
您可以通过使用 mlx-image 库名称进行筛选来查找 mlx-image 模型，例如 此查询。
此外，还有一个开放的 mlx-vision 社区，供贡献者转换并发布 MLX 格式的权重。
安装
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 结果
前往 results-imagenet-1k.csv 查看每个转换为 mlx-image 的模型及其在不同设置下于 ImageNet-1K 上的性能。
TL;DR 性能与 PyTorch 实现的原始模型相当。
与 PyTorch 及其他熟悉工具的相似性
mlx-image 力求尽可能贴近 PyTorch：
DataLoader-> 您可以自定义collate_fn，并使用num_workers来加速数据加载Dataset->mlx-image已经支持LabelFolderDataset（经典 PyTorchImageFolder）和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。