apple-aiml-research/ml-mobileone

What is MobileOne?

MobileOne is a family of ultra‑fast image‑classification backbones designed for mobile devices. The repository provides the PyTorch implementation of the models described in the paper “An Improved One millisecond Mobile Backbone” (arXiv 2206.04040) together with pretrained weights and a tiny iOS app for measuring real‑world latency.

---

Key points

Goal	Deliver high‑accuracy image classifiers that run in ≤ 1 ms on a phone (iPhone 12 Pro).
Models	Five variants – MobileOne‑S0 … S4 – trading off accuracy (71.4 % → 79.4 % Top‑1 on ImageNet) for latency (0.79 ms → 1.86 ms).
How it works	During training the network contains parallel “branch” convolutions that improve optimization. At inference time a re‑parameterization step fuses those branches into a single, standard convolution, giving the speed of a plain network.
Artifacts	• PyTorch checkpoints (fused & unfused)
• CoreML models for iOS
• A small iOS benchmark app (ModelBench) to verify latency on real hardware.
Typical use‑cases	• Fine‑tuning on a custom image dataset for mobile deployment.
• Direct inference on‑device via CoreML.
• Research on fast CNN architectures.

---

Getting started (code snippet from the README)

import torch
from mobileone import mobileone, reparameterize_model

# 1️⃣ Build a model (choose a variant)
model = mobileone(variant='s0')          # or 's1' … 's4'

# 2️⃣ Train from scratch or fine‑tune …
# model = … (your training loop)

# 3️⃣ Load a pretrained checkpoint for further training
ckpt = torch.load('/path/to/unfused_checkpoint.pth.tar')
model.load_state_dict(ckpt)

# 4️⃣ Prepare for fast inference
model.eval()
model_eval = reparameterize_model(model)   # fuses branches
# → use `model_eval` for validation or deployment

For pure evaluation you can skip the re‑parameterization step by loading the fused checkpoint directly:

model = mobileone(variant='s0', inference_mode=True)
model.load_state_dict(torch.load('/path/to/checkpoint.pth.tar'))
# ready to run inference

---

Where to find the models

Variant	Top‑1 Acc.	Latency* (iPhone 12 Pro)	PyTorch checkpoint	CoreML model
S0	71.4 %	0.79 ms	download (fused) / unfused	mlmodel
S1	75.9 %	0.89 ms	…	…
S2	77.4 %	1.18 ms	…	…
S3	78.1 %	1.53 ms	…	…
S4	79.4 %	1.86 ms	…	…

*Latency measured on an iPhone 12 Pro.

---

iOS benchmark app (ModelBench)

The ModelBench folder contains a minimal Xcode project that loads the CoreML version of a MobileOne model and reports the actual inference time on the device. The README inside that folder explains how to build and run the app.

---

Citing the work

If you use MobileOne in research or a product, cite the original paper:

@article{mobileone2022,
  title={An Improved One millisecond Mobile Backbone},
  author={Vasu, Pavan Kumar Anasosalu and Gabriel, James and Zhu, Jeff and Tuzel, Oncel and Ranjan, Anurag},
  journal={arXiv preprint arXiv:2206.04040},
  year={2022}
}

---

TL;DR

MobileOne is a ready‑to‑use, PyTorch‑based suite of mobile‑friendly CNNs that achieve state‑of‑the‑art ImageNet accuracy while staying under a millisecond per image on modern iPhones. The repo ships training code, pretrained weights (both fused for inference and unfused for further training), CoreML exports, and a tiny iOS benchmark app.