"""Train the same two-layer MNIST network with MLX, Apple's array framework. Purpose: the Track M native path for the first training run, kept line for line as close to train-mnist.py as the two frameworks allow, so the differences are visible: arrays live in unified memory, computation is lazy until mx.eval, and gradients come from a function transform (nn.value_and_grad) rather than from .backward(). Platform: mac (Apple silicon, macOS 14 or later, native arm64 Python, mlx installed); also runs on the CPU build of MLX on Linux (pip install "mlx[cpu]") for checking. Minimum memory: 8 GB Assumes: mlx and torchvision are installed in the active environment (torchvision is used only to download and decode MNIST), and about 70 MB of free disk under ./data. Usage: python train-mnist-mlx.py [--epochs 5] [--lr 0.1] [--batch-size 128] [--hidden 256] [--train-size 50000] [--device gpu|cpu] [--checkpoint mnist_mlp.safetensors] [--label mlx-first-run] [--labbook labbook.md] [--seed 0] [--data ./data] The output format is the same as train-mnist.py so the two logs can be compared column by column; the checkpoint is a .safetensors file, which MLX writes natively. """ import argparse import json import math import sys import time from pathlib import Path import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim import numpy as np from mlx.utils import tree_flatten from torchvision import datasets def load_split(root: str, train: bool) -> tuple[mx.array, mx.array]: ds = datasets.MNIST(root, train=train, download=True) images = ds.data.numpy().astype("float32") / 255.0 # bytes 0..255 -> 0.0..1.0 labels = ds.targets.numpy().astype("int32") return mx.array(images.reshape(len(images), -1)), mx.array(labels) # (n, 784), (n,) class TwoLayerNet(nn.Module): """784 inputs -> hidden units with ReLU -> 10 outputs, one per digit.""" def __init__(self, hidden: int = 256) -> None: super().__init__() self.l1 = nn.Linear(28 * 28, hidden) self.l2 = nn.Linear(hidden, 10) def __call__(self, x: mx.array) -> mx.array: return self.l2(nn.relu(self.l1(x))) def loss_fn(model: TwoLayerNet, x: mx.array, y: mx.array) -> mx.array: return mx.mean(nn.losses.cross_entropy(model(x), y)) def evaluate(model: TwoLayerNet, x: mx.array, y: mx.array) -> tuple[float, float]: logits = model(x) loss = mx.mean(nn.losses.cross_entropy(logits, y)).item() acc = mx.mean(mx.argmax(logits, axis=1) == y).item() return loss, acc def main() -> None: parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--epochs", type=int, default=5) parser.add_argument("--lr", type=float, default=0.1, help="learning rate for plain SGD") parser.add_argument("--batch-size", type=int, default=128) parser.add_argument("--hidden", type=int, default=256) parser.add_argument("--train-size", type=int, default=50_000, help="how many of the 50,000 training images to train on") parser.add_argument("--device", default="gpu", choices=["gpu", "cpu"], help="MLX default device; gpu is the Apple GPU through Metal") parser.add_argument("--checkpoint", default="mnist_mlp.safetensors") parser.add_argument("--label", default="", help="a short name for this run in the notebook") parser.add_argument("--labbook", default=None, help="append one JSON line per run to this file") parser.add_argument("--seed", type=int, default=0) parser.add_argument("--data", default="./data") args = parser.parse_args() if not 1 <= args.train_size <= 50_000: sys.exit("train-mnist-mlx: --train-size must be between 1 and 50000") if not args.checkpoint.endswith((".safetensors", ".npz")): sys.exit("train-mnist-mlx: --checkpoint must end in .safetensors or .npz (the formats save_weights writes)") if args.device == "cpu": mx.set_default_device(mx.cpu) elif mx.default_device() != mx.gpu: sys.exit(f"train-mnist-mlx: the default device is {mx.default_device()}, not the GPU; " "pass --device cpu on a build without Metal") mx.random.seed(args.seed) np.random.seed(args.seed) print(f"mlx {mx.__version__}") print(f"device: {mx.default_device()}") x_all, y_all = load_split(args.data, train=True) x_test, y_test = load_split(args.data, train=False) x_train, y_train = x_all[:args.train_size], y_all[:args.train_size] # the first 50,000 (or fewer) x_val, y_val = x_all[50_000:], y_all[50_000:] # the last 10,000 held out n = x_train.shape[0] steps = math.ceil(n / args.batch_size) print(f"data: train {n:,} val {x_val.shape[0]:,} test {x_test.shape[0]:,} " f"batch {args.batch_size} -> {steps} steps per epoch") model = TwoLayerNet(args.hidden) mx.eval(model.parameters()) # parameters are lazy until evaluated n_params = sum(p.size for _, p in tree_flatten(model.parameters())) print(f"parameters: {n_params:,}") for name, p in tree_flatten(model.parameters()): print(f" {name:20s} shape {tuple(p.shape)}") optimiser = optim.SGD(learning_rate=args.lr) step = nn.value_and_grad(model, loss_fn) # a function returning (loss, grads) for the model's parameters val_loss, val_acc = evaluate(model, x_val, y_val) print(f"epoch 0 (untrained) val loss {val_loss:.4f} val acc {val_acc:.4f}") history, best_val, best_epoch = [], float("inf"), 0 started = time.time() for epoch in range(1, args.epochs + 1): epoch_started = time.time() order = np.random.permutation(n) running = 0.0 for start in range(0, n, args.batch_size): idx = mx.array(order[start:start + args.batch_size]) loss, grads = step(model, x_train[idx], y_train[idx]) # forward + loss + backward optimiser.update(model, grads) # step mx.eval(model.parameters(), optimiser.state) # force the lazy graph to run running += loss.item() * idx.shape[0] train_loss = running / n val_loss, val_acc = evaluate(model, x_val, y_val) seconds = time.time() - epoch_started history.append({"epoch": epoch, "train_loss": round(train_loss, 4), "val_loss": round(val_loss, 4), "val_acc": round(val_acc, 4), "seconds": round(seconds, 1)}) marker = "" if val_loss < best_val: best_val, best_epoch = val_loss, epoch model.save_weights(args.checkpoint) marker = " (saved)" print(f"epoch {epoch:2d} train loss {train_loss:.4f} val loss {val_loss:.4f} " f"val acc {val_acc:.4f} {seconds:5.1f} s{marker}") elapsed = time.time() - started if best_epoch == 0: test_loss, test_acc = float("nan"), float("nan") print(f"no epoch improved the validation loss: nothing saved to {args.checkpoint}, " f"test set not used ({elapsed:.0f} s)") else: model.load_weights(args.checkpoint) test_loss, test_acc = evaluate(model, x_test, y_test) print(f"best epoch {best_epoch}: test loss {test_loss:.4f} test acc {test_acc:.4f} " f"({elapsed:.0f} s total, checkpoint {args.checkpoint})") if args.labbook: record = { "lab": "part-01/train-mnist-mlx", "label": args.label, "device": str(mx.default_device()), "mlx": mx.__version__, "epochs": args.epochs, "lr": args.lr, "batch_size": args.batch_size, "hidden": args.hidden, "train_size": n, "parameters": n_params, "best_epoch": best_epoch, "test_acc": None if math.isnan(test_acc) else round(test_acc, 4), "seconds": round(elapsed, 1), "checkpoint": args.checkpoint, "history": history, } with Path(args.labbook).open("a", encoding="utf-8") as fh: fh.write(json.dumps(record) + "\n") print(f"recorded in {args.labbook}") if __name__ == "__main__": main()