diff --git a/AGENTS.md b/AGENTS.md
index 853a4de..10efe91 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -39,7 +39,7 @@ that lazily converts observations to Julia types).
`from_csv`/`from_json`/`from_parquet`) that re-wrap results in the default
`"julia"` format.
- `src/transforms.jl` — `py2jl` / `numpy2jl` / `jl2numpy` / `jl2py`. `py2jl`
- recursively converts Python containers, numpy arrays (copyless via DLPack), and
+ recursively converts Python containers, numpy arrays (copyless, zero-copy), and
PIL images into Julia types; `jl2py` is the write-path dual. The `"julia"`
format is numpy-backed, so numeric array columns decode to real N-D Julia arrays
and image columns decode to raw numeric arrays (not `Colorant` colorviews).
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 37d612c..e498332 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
## [Unreleased]
+### Changed
+- Dropped the `DLPack` dependency. `numpy2jl` (the zero-copy numpy → Julia conversion behind
+ `py2jl`) now shares the buffer through PythonCall directly (`PyArray` + `unsafe_wrap`), still
+ returning a genuine `Array` so results stay on Julia's fast paths (BLAS, GPU host→device copies).
+ Behavior is unchanged — zero-copy, write-back, reversed axes; read-only/non-contiguous buffers are
+ still copied. A bonus fix: buffer cleanup now routes through PythonCall's GIL-deferred decref
+ instead of DLPack's GIL-reentering finalizer, so a numpy buffer freed on a `DataLoader`
+ `parallel=true` worker thread can no longer deadlock against a thread compiling under the GIL.
+
### Added
- `MLUtils.DataLoader(ds::Dataset; num_workers=N)` works out of the box for process-parallel
loading, including over `mapobs`/`ObsView`-wrapped datasets. Building on the `Serialization`
diff --git a/Project.toml b/Project.toml
index 153911b..93879d5 100644
--- a/Project.toml
+++ b/Project.toml
@@ -9,7 +9,6 @@ projects = ["test", "docs"]
[deps]
Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
-DLPack = "53c2dc0f-f7d5-43fd-8906-6c0220547083"
ImageCore = "a09fc81d-aa75-5fe9-8630-4744c3626534"
MLCore = "c2834f40-e789-41da-a90e-33b280584a8c"
PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
@@ -19,7 +18,6 @@ Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
[compat]
Compat = "4.10"
CondaPkg = "0.2"
-DLPack = "0.3"
ImageCore = "0.9, 0.10"
MLCore = "1.1"
PythonCall = "0.9"
diff --git a/README.md b/README.md
index e6aee86..0ad1e61 100644
--- a/README.md
+++ b/README.md
@@ -26,7 +26,7 @@ pkg> add HuggingFaceDatasets
HuggingFaceDatasets.jl provides wrappers around types from the `datasets` python package (e.g. `Dataset` and `DatasetDict`) along with a few related methods.
-Check out the [examples/](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/tree/main/examples) folder for usage examples.
+Check out the [perf/](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/tree/main/perf) folder for usage examples and data-loading benchmarks.
Observations are returned in the `"julia"` format by default, i.e. converted to native julia types on access:
diff --git a/docs/src/guide.md b/docs/src/guide.md
index fc152d0..5998bfe 100644
--- a/docs/src/guide.md
+++ b/docs/src/guide.md
@@ -293,7 +293,7 @@ Note that the format and the custom transform share the same slot: setting the `
format installs `py2jl` as the transform, and `with_jltransform` then replaces it — which
is why the transform above calls `py2jl` itself. For layering additional per-batch
processing on top of the `"julia"` format, prefer `MLUtils.mapobs` (see below), as in the
-[`examples/flux_mnist/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/examples/flux_mnist/flux_mnist.jl)
+[`perf/flux_mnist/flux_mnist.jl`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/blob/main/perf/flux_mnist/flux_mnist.jl)
script.
The order of operations when you index is:
@@ -367,7 +367,7 @@ A few things to note:
!!! warning "Arrays come back transposed"
numpy is **row-major**, Julia is **column-major**. The zero-copy conversion
- ([`numpy2jl`](@ref), via DLPack) therefore returns an array whose **dimensions are
+ ([`numpy2jl`](@ref)) therefore returns an array whose **dimensions are
reversed** relative to the Python side: a numpy array of shape `(d₁, …, dₙ)` becomes a
Julia array of size `(dₙ, …, d₁)`.
@@ -598,4 +598,4 @@ guide for the mechanics and tradeoffs.
A complete example — `mapobs` transform, `num_workers`, and a training loop, benchmarked
against PyTorch — lives in
-[`examples/flux_mnist/`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/tree/main/examples/flux_mnist).
+[`perf/flux_mnist/`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/tree/main/perf/flux_mnist).
diff --git a/docs/src/index.md b/docs/src/index.md
index 20b1902..2cd1286 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -86,7 +86,7 @@ Python: {'label': 5}
See the [Guide](@ref) for the transform workflow, method forwarding, array/image
orientation, and integration with MLUtils/Flux data loaders. Runnable examples
live in the
-[`examples/`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/tree/main/examples)
+[`perf/`](https://github.com/JuliaGenAI/HuggingFaceDatasets.jl/tree/main/perf)
folder.
## Troubleshooting
diff --git a/examples/flux_mnist/README.md b/examples/flux_mnist/README.md
deleted file mode 100644
index 63af6db..0000000
--- a/examples/flux_mnist/README.md
+++ /dev/null
@@ -1,80 +0,0 @@
-# Flux MNIST — data-loading benchmark
-
-Trains a small MLP on MNIST pulled from the HuggingFace `datasets` library, comparing
-**data-loading strategies** (on-the-fly vs. materialized, thread- vs. process-parallel) in
-Flux+HuggingFaceDatasets.jl against equivalent and idiomatic PyTorch versions.
-
-## Files
-
-| file | stack | notes |
-| --- | --- | --- |
-| [`flux_mnist.jl`](flux_mnist.jl) | Flux + HuggingFaceDatasets.jl | `julia --project=. -t4 flux_mnist.jl` |
-| [`pytorch_mnist.py`](pytorch_mnist.py) | PyTorch + `datasets` | 1:1 port of `flux_mnist.jl`; `uv run pytorch_mnist.py` |
-| [`pytorch_mnist_idiomatic.py`](pytorch_mnist_idiomatic.py) | PyTorch + `datasets` | idiomatic HF patterns (`with_transform`, `.map`); `uv run pytorch_mnist_idiomatic.py` |
-
-## Setup
-
-Same model and hyperparameters everywhere: MLP `784 → 100 → 100 → 10`, `AdamW(1e-3)`,
-cross-entropy, batch size 128, **4 epochs**, **CPU only**, MNIST (`ylecun/mnist`) via HF
-`datasets`. Timings are single-run wall-clock on an **Apple M1 Pro** — indicative, not
-rigorous (expect ±10–20% run to run). A warm-up epoch precedes the Julia timings to exclude
-Julia's one-time JIT compilation; PyTorch runs eagerly (no comparable compile step) so needs
-none. Both then measure steady-state per-epoch compute — imports and process startup sit
-outside the timed region either way — so the tables are directly comparable.
-
-## Results
-
-### Julia — Flux + HuggingFaceDatasets.jl (`-t4`)
-
-| config | data loading | time |
-| --- | --- | ---: |
-| Serial | on-the-fly, `num_workers=0` | 13.5 s |
-| Serial Materialized | in-memory `[:]`, `num_workers=0` | 6.4 s |
-| Parallel Materialized | in-memory, `parallel=true` (threads) | 9.1 s |
-| Distributed | on-the-fly, `num_workers=4` (processes) | 33.0 s |
-
-A warm-up epoch runs first, so these exclude Julia's JIT compilation. The serial path is
-then fully warmed; the `parallel`/`num_workers` paths still compile their first call, and
-Distributed also pays worker-process startup — real costs for a short job, left in the numbers.
-
-### PyTorch — 1:1 port (`pytorch_mnist.py`)
-
-| config | data loading | time |
-| --- | --- | ---: |
-| Serial | on-the-fly, `num_workers=0` | 62.8 s |
-| Serial Materialized | in-memory tensors, `num_workers=0` | 16.8 s |
-| Parallel Materialized | in-memory, `num_workers=4` (processes) | 22.2 s |
-| Distributed | on-the-fly, `num_workers=4` (processes) | 66.6 s |
-
-### PyTorch — idiomatic HF (`pytorch_mnist_idiomatic.py`)
-
-| config | data loading | time |
-| --- | --- | ---: |
-| Lazy `with_transform` | on-the-fly, `num_workers=0` | 34.1 s |
-| Lazy `with_transform` | on-the-fly, `num_workers=4` | 55.1 s |
-| Cached `.map` (torch format) | preprocessed to Arrow, `num_workers=0` | 291.4 s† |
-| Cached `.map` (torch format) | preprocessed to Arrow, `num_workers=4` | 128.3 s |
-
-†includes the one-time `.map` decode+cache; the `num_workers=4` row reuses that cache.
-
-## Takeaways
-
-- **Materializing into memory is the biggest win** for a dataset this small — ~2× in Julia
- (13.5 → 6.4 s) and ~4× in PyTorch (62.8 → 16.8 s), because it drops the per-batch CPython
- decode entirely.
-- **Parallel loading does *not* pay off here.** The MLP is tiny, so multiprocess
- pickling/IPC overhead outweighs the parallel-decode benefit — `num_workers=4` was *slower*
- than `num_workers=0` in every PyTorch case, and Julia's `Distributed` (33.0 s) lost to its
- serial on-the-fly (13.5 s). `num_workers`/process parallelism is for when `getobs` is the
- bottleneck (large images, heavy decode), not toy workloads.
- - Julia's **thread**-based `parallel=true` over *materialized* data (9.1 s) is cheap
- (shared memory, needs `-t>1`) but pointless here — there's nothing to parallelize once the
- data is plain in-memory arrays. PyTorch has no thread analog (the GIL); its only knob is
- `num_workers` (processes).
-- **Idiomatic PyTorch is competitive but has a trap.** Lazy `with_transform` (34.1 s) even
- beats the 1:1 Serial port (62.8 s), but naively `.map`-ing decoded float tensors into Arrow
- and reading them back via the torch formatter is *much* slower (128–291 s) — for images,
- prefer lazy transforms or materialize to plain tensors, don't cache decoded floats.
-- **Julia came out ahead of PyTorch** on the like-for-like paths (on-the-fly 13.5 vs 62.8 s;
- materialized 6.4 vs 16.8 s) — a fair comparison, since both are steady-state per-epoch
- compute (Julia after a warm-up, PyTorch eager with no compile step to warm away).
diff --git a/examples/flux_mnist/flux_mnist.jl b/examples/flux_mnist/flux_mnist.jl
deleted file mode 100644
index 874e9e5..0000000
--- a/examples/flux_mnist/flux_mnist.jl
+++ /dev/null
@@ -1,93 +0,0 @@
-using Random, Statistics
-using Flux
-using Flux.Losses: logitcrossentropy
-using Flux: onecold, onehotbatch
-using HuggingFaceDatasets
-using MLUtils: MLUtils, mapobs
-# using ProfileView, BenchmarkTools
-
-function mnist_transform(batch)
- # the image column is a stacked (W, H, N) UInt8 array,
- # so just rescale to Float32 in [0, 1]
- image = batch["image"] ./ 255f0
- return (; image, label = batch["label"])
-end
-
-function loss_and_accuracy(data_loader, model, device)
- acc = 0
- ls = 0.0f0
- num = 0
- for (x, y) in data_loader
- x = x |> device
- yoh = onehotbatch(y, 0:9) |> device
- ŷ = model(x)
- ls += logitcrossentropy(ŷ, yoh, agg=sum)
- acc += sum(onecold(ŷ, 0:9) .== y)
- num += length(y)
- end
- return ls / num, acc / num
-end
-
-# `num_workers = 0` loads on the main process; `num_workers > 0` spreads each batch's
-# `getobs` (and the CPython read it triggers) over that many worker processes, sidestepping
-# the GIL. MLUtils spawns the workers on demand under the current `--project`.
-function train(; epochs=2, num_workers=0, materialize=false, parallel=false, verbose=true)
- batchsize = 128
- nhidden = 100
- device = cpu
-
- train_data = load_dataset("ylecun/mnist", split="train")
- test_data = load_dataset("ylecun/mnist", split="test")
- # apply the transform lazily so it runs per batch during iteration (on the workers when
- # `num_workers > 0`); `mapobs`/`ObsView`-wrapped datasets compose with `num_workers`
- train_data = mapobs(mnist_transform, train_data)
- test_data = mapobs(mnist_transform, test_data)
- if materialize
- train_data = train_data[:]
- test_data = test_data[:]
- end
-
- train_loader = Flux.DataLoader(train_data; batchsize, shuffle=true, num_workers, parallel)
- test_loader = Flux.DataLoader(test_data; batchsize, num_workers, parallel)
-
- model = Chain([Flux.flatten,
- Dense(28*28, nhidden, relu),
- Dense(nhidden, nhidden, relu),
- Dense(nhidden, 10)]) |> device
-
- opt = Flux.setup(AdamW(1e-3), model)
-
- function report(epoch)
- train_loss, train_acc = loss_and_accuracy(train_loader, model, device)
- test_loss, test_acc = loss_and_accuracy(test_loader, model, device)
- r(x) = round(x, digits=3)
- r(x::Int) = x
- @info map(r, (; epoch, train_loss, train_acc, test_loss, test_acc))
- end
-
- verbose && report(0)
- for epoch in 1:epochs
- for (x, y) in train_loader
- x = x |> device
- yoh = onehotbatch(y, 0:9) |> device
- loss, grads = Flux.withgradient(m -> logitcrossentropy(m(x), yoh), model)
- Flux.update!(opt, model, grads[1])
- end
- verbose && report(epoch)
- end
-end
-
-println("#### START COMPARISON ###########")
-MLUtils.close_dataloader_pool()
-println("### WARMUP") # for precompilation
-@time train(; epochs=1, num_workers=0, materialize=false, verbose=false)
-println("### Serial")
-@time train(; epochs=4, num_workers=0, materialize=false, verbose=false)
-println("### Serial Materialized")
-@time train(; epochs=4, num_workers=0, materialize=true, verbose=false)
-println("### Paralle Materialized")
-@time train(; epochs=4, num_workers=0, materialize=true, parallel=true, verbose=false)
-println("### Distributed")
-@time train(; epochs=4, num_workers=4, materialize=false, verbose=false)
-MLUtils.close_dataloader_pool()
-println("#### END COMPARISON ###########")
diff --git a/examples/flux_mnist/pytorch_mnist.py b/examples/flux_mnist/pytorch_mnist.py
deleted file mode 100644
index 9f4d910..0000000
--- a/examples/flux_mnist/pytorch_mnist.py
+++ /dev/null
@@ -1,135 +0,0 @@
-# /// script
-# requires-python = ">=3.10"
-# dependencies = ["torch", "datasets", "numpy", "pillow"]
-# ///
-#
-# PyTorch counterpart of flux_mnist.jl (Flux + HuggingFaceDatasets), for a data-loading timing
-# comparison on the same task: an MLP on MNIST pulled from the same HuggingFace `datasets`
-# Arrow dataset. Run with: uv run examples/flux_mnist/pytorch_mnist.py
-#
-# The four configs mirror flux_mnist.jl. One caveat: PyTorch's only DataLoader parallelism is
-# multiprocess (`num_workers`), so the "Parallel Materialized" row uses worker *processes*
-# over in-memory tensors — Julia's `parallel=true` there uses *threads* (cheap, shared
-# memory), which PyTorch cannot do because of the GIL. This is a timing benchmark, not a
-# numerical match: the exact AdamW weight decay etc. differ, so accuracies only track loosely.
-
-import time
-
-import numpy as np
-import torch
-import torch.nn as nn
-from datasets import load_dataset, disable_progress_bars
-from datasets.utils.logging import set_verbosity_error
-from torch.utils.data import DataLoader, Dataset, TensorDataset
-
-disable_progress_bars()
-set_verbosity_error()
-
-BATCHSIZE = 128
-NHIDDEN = 100
-EPOCHS = 4
-DEVICE = torch.device("cpu")
-
-
-def make_model():
- return nn.Sequential(
- nn.Flatten(),
- nn.Linear(28 * 28, NHIDDEN), nn.ReLU(),
- nn.Linear(NHIDDEN, NHIDDEN), nn.ReLU(),
- nn.Linear(NHIDDEN, 10),
- ).to(DEVICE)
-
-
-class HFImageDataset(Dataset):
- """On-the-fly loading: index the Arrow-backed HF dataset and decode per item.
-
- Mirrors `mapobs(mnist_transform, ds)` in flux_mnist.jl. The dataset is memory-mapped and
- picklable, so with num_workers>0 each worker process re-opens it (data shared, not
- copied) — the same mechanism as the Julia `num_workers` path.
- """
-
- def __init__(self, hf_ds):
- self.ds = hf_ds.with_format("numpy")
-
- def __len__(self):
- return len(self.ds)
-
- def __getitem__(self, i):
- row = self.ds[i]
- image = np.asarray(row["image"], dtype=np.float32) / 255.0 # (28, 28) in [0, 1]
- return torch.from_numpy(image), int(row["label"])
-
-
-def materialize(hf_ds):
- """Decode the whole split into in-memory tensors up front (the `[:]` path in flux_mnist.jl)."""
- ds = hf_ds.with_format("numpy")
- col = ds["image"]
- images = col if isinstance(col, np.ndarray) else np.stack(col)
- images = torch.from_numpy(images.astype(np.float32) / 255.0) # (N, 28, 28)
- labels = torch.from_numpy(np.asarray(ds["label"]).astype(np.int64)) # (N,)
- return TensorDataset(images, labels)
-
-
-@torch.no_grad()
-def loss_and_accuracy(loader, model):
- model.eval()
- lossfn = nn.CrossEntropyLoss(reduction="sum")
- ls, correct, num = 0.0, 0, 0
- for x, y in loader:
- x, y = x.to(DEVICE), y.to(DEVICE)
- logits = model(x)
- ls += lossfn(logits, y).item()
- correct += (logits.argmax(1) == y).sum().item()
- num += y.shape[0]
- return ls / num, correct / num
-
-
-def train(num_workers=0, materialize_data=False):
- train_hf = load_dataset("ylecun/mnist", split="train")
- test_hf = load_dataset("ylecun/mnist", split="test")
-
- if materialize_data:
- train_ds, test_ds = materialize(train_hf), materialize(test_hf)
- else:
- train_ds, test_ds = HFImageDataset(train_hf), HFImageDataset(test_hf)
-
- # persistent_workers keeps the pool alive across epochs (closer to Julia's leased pool);
- # without it PyTorch respawns workers on every pass over the loader.
- kw = dict(num_workers=num_workers, persistent_workers=num_workers > 0)
- train_loader = DataLoader(train_ds, batch_size=BATCHSIZE, shuffle=True, **kw)
- test_loader = DataLoader(test_ds, batch_size=BATCHSIZE, **kw)
-
- model = make_model()
- opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
- lossfn = nn.CrossEntropyLoss()
-
- def report(epoch):
- tr_loss, tr_acc = loss_and_accuracy(train_loader, model)
- te_loss, te_acc = loss_and_accuracy(test_loader, model)
- print(f"(epoch = {epoch}, train_loss = {tr_loss:.3f}, train_acc = {tr_acc:.3f}, "
- f"test_loss = {te_loss:.3f}, test_acc = {te_acc:.3f})")
-
- report(0)
- for epoch in range(1, EPOCHS + 1):
- model.train()
- for x, y in train_loader:
- x, y = x.to(DEVICE), y.to(DEVICE)
- opt.zero_grad()
- loss = lossfn(model(x), y)
- loss.backward()
- opt.step()
- report(epoch)
-
-
-def timed(name, **kw):
- print(f"### {name}")
- t0 = time.perf_counter()
- train(**kw)
- print(f" {time.perf_counter() - t0:.3f} seconds")
-
-
-if __name__ == "__main__":
- timed("Serial", num_workers=0, materialize_data=False)
- timed("Serial Materialized", num_workers=0, materialize_data=True)
- timed("Parallel Materialized", num_workers=4, materialize_data=True)
- timed("Distributed", num_workers=4, materialize_data=False)
diff --git a/examples/flux_mnist/pytorch_mnist_idiomatic.py b/examples/flux_mnist/pytorch_mnist_idiomatic.py
deleted file mode 100644
index ae8057b..0000000
--- a/examples/flux_mnist/pytorch_mnist_idiomatic.py
+++ /dev/null
@@ -1,151 +0,0 @@
-# /// script
-# requires-python = ">=3.10"
-# dependencies = ["torch", "torchvision", "datasets", "numpy", "pillow", "tqdm"]
-# ///
-#
-# Idiomatic HuggingFace `datasets` + PyTorch MNIST training — written the way a practitioner
-# typically would, rather than line-for-line equivalent to the Flux `flux_mnist.jl` (see
-# `pytorch_mnist.py` for that). Run with: uv run examples/flux_mnist/pytorch_mnist_idiomatic.py
-#
-# Idioms shown:
-# * `load_dataset` -> DatasetDict, then `ds["train"]` / `ds["test"]`
-# * a torchvision transform applied lazily via `Dataset.with_transform` (the canonical image
-# pipeline) — plus the `.map(batched=True)` + torch-format cached alternative
-# * a `collate_fn` feeding a standard `DataLoader(num_workers=...)`
-# * a plain `nn.Module`, `AdamW`, and a `tqdm` training loop
-#
-# The model, batch size and epochs match the Flux run so data-loading timings are comparable.
-# Kept on CPU for that comparison; a typical script would auto-select the accelerator (see below).
-
-import time
-
-import numpy as np
-import torch
-import torch.nn as nn
-from datasets import load_dataset, disable_progress_bars
-from datasets.utils.logging import set_verbosity_error
-from torch.utils.data import DataLoader
-from torchvision.transforms import Compose, Normalize, ToTensor
-from tqdm import tqdm
-
-disable_progress_bars()
-set_verbosity_error()
-
-BATCHSIZE = 128
-EPOCHS = 4
-# A typical script auto-selects the accelerator:
-# DEVICE = torch.device("cuda" if torch.cuda.is_available()
-# else "mps" if torch.backends.mps.is_available() else "cpu")
-# Pinned to CPU here so the timings line up with the Flux example.
-DEVICE = torch.device("cpu")
-
-# ToTensor scales uint8 [0, 255] -> float [0, 1] and lays out (C, H, W); Normalize applies the
-# standard MNIST mean/std. This is the usual torchvision image pipeline.
-IMAGE_TF = Compose([ToTensor(), Normalize((0.1307,), (0.3081,))])
-
-
-class MLP(nn.Module):
- def __init__(self, nhidden=100):
- super().__init__()
- self.net = nn.Sequential(
- nn.Flatten(),
- nn.Linear(28 * 28, nhidden), nn.ReLU(),
- nn.Linear(nhidden, nhidden), nn.ReLU(),
- nn.Linear(nhidden, 10),
- )
-
- def forward(self, x):
- return self.net(x)
-
-
-def lazy_transform(batch):
- """Decode + transform on access (registered with `with_transform`)."""
- batch["pixel_values"] = [IMAGE_TF(img) for img in batch["image"]]
- return batch
-
-
-def collate(examples):
- x = torch.stack([e["pixel_values"] for e in examples])
- y = torch.tensor([e["label"] for e in examples])
- return x, y
-
-
-def cached_preprocess(batch):
- """Same transform, but returning arrays so `datasets` can cache them to Arrow."""
- batch["pixel_values"] = [IMAGE_TF(img).numpy() for img in batch["image"]]
- return batch
-
-
-def make_loaders(num_workers, cached):
- ds = load_dataset("ylecun/mnist")
- if cached:
- # Preprocess once; `datasets` caches the result to Arrow and hands back torch tensors
- # on access, so the default collate works and no per-item Python decode is needed.
- ds = ds.map(cached_preprocess, batched=True, remove_columns=["image"])
- ds = ds.with_format("torch")
- train_split, test_split, collate_fn = ds["train"], ds["test"], None
- else:
- train_split = ds["train"].with_transform(lazy_transform)
- test_split = ds["test"].with_transform(lazy_transform)
- collate_fn = collate
-
- kw = dict(num_workers=num_workers, persistent_workers=num_workers > 0, collate_fn=collate_fn)
- train_loader = DataLoader(train_split, batch_size=BATCHSIZE, shuffle=True, **kw)
- test_loader = DataLoader(test_split, batch_size=BATCHSIZE, **kw)
- return train_loader, test_loader
-
-
-def _xy(batch):
- # lazy path yields (x, y) tuples; cached torch-format path yields column dicts.
- if isinstance(batch, dict):
- return batch["pixel_values"], batch["label"]
- return batch
-
-
-@torch.no_grad()
-def evaluate(loader, model, lossfn):
- model.eval()
- total_loss, correct, num = 0.0, 0, 0
- for batch in loader:
- x, y = _xy(batch)
- x, y = x.to(DEVICE), y.to(DEVICE)
- logits = model(x)
- total_loss += lossfn(logits, y).item() * y.shape[0]
- correct += (logits.argmax(1) == y).sum().item()
- num += y.shape[0]
- return total_loss / num, correct / num
-
-
-def train(num_workers=0, cached=False):
- train_loader, test_loader = make_loaders(num_workers, cached)
- model = MLP().to(DEVICE)
- opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
- lossfn = nn.CrossEntropyLoss()
-
- for epoch in range(1, EPOCHS + 1):
- model.train()
- for batch in tqdm(train_loader, desc=f"epoch {epoch}/{EPOCHS}", leave=False):
- x, y = _xy(batch)
- x, y = x.to(DEVICE), y.to(DEVICE)
- opt.zero_grad()
- loss = lossfn(model(x), y)
- loss.backward()
- opt.step()
- tr_loss, tr_acc = evaluate(train_loader, model, lossfn)
- te_loss, te_acc = evaluate(test_loader, model, lossfn)
- print(f"(epoch = {epoch}, train_loss = {tr_loss:.3f}, train_acc = {tr_acc:.3f}, "
- f"test_loss = {te_loss:.3f}, test_acc = {te_acc:.3f})")
-
-
-def timed(name, **kw):
- print(f"### {name}")
- t0 = time.perf_counter()
- train(**kw)
- print(f" {time.perf_counter() - t0:.3f} seconds")
-
-
-if __name__ == "__main__":
- timed("Lazy with_transform, num_workers=0", cached=False, num_workers=0)
- timed("Lazy with_transform, num_workers=4", cached=False, num_workers=4)
- timed("Cached .map (torch format), num_workers=0", cached=True, num_workers=0)
- timed("Cached .map (torch format), num_workers=4", cached=True, num_workers=4)
diff --git a/perf/flux_cifar10/Project.toml b/perf/flux_cifar10/Project.toml
new file mode 100644
index 0000000..d334b29
--- /dev/null
+++ b/perf/flux_cifar10/Project.toml
@@ -0,0 +1,14 @@
+[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+HuggingFaceDatasets = "d94b9a45-fdf5-4270-b024-5cbb9ef7117d"
+MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
+
+[sources]
+HuggingFaceDatasets = {path = "../.."}
+
+[compat]
+MLUtils = "0.4.12"
diff --git a/perf/flux_cifar10/README.md b/perf/flux_cifar10/README.md
new file mode 100644
index 0000000..5227f88
--- /dev/null
+++ b/perf/flux_cifar10/README.md
@@ -0,0 +1,94 @@
+# Flux CIFAR-10 — CNN + data-loading benchmark (GPU)
+
+Trains a small VGG-style CNN on CIFAR-10 pulled from the HuggingFace `datasets` library, with the
+standard crop+flip augmentation, **on the GPU**. It compares **data-loading strategies** (on-the-fly
+vs. materialized, thread- vs. process-parallel) in Flux+HuggingFaceDatasets.jl against an equivalent
+PyTorch version — the CIFAR-10 counterpart of the [`flux_mnist`](../flux_mnist) example.
+
+## Files
+
+| file | stack | notes |
+| --- | --- | --- |
+| [`flux_cifar10.jl`](flux_cifar10.jl) | Flux + HuggingFaceDatasets.jl | `julia --project=. -t4 flux_cifar10.jl` |
+| [`pytorch_cifar10.py`](pytorch_cifar10.py) | PyTorch + `datasets` | `uv run pytorch_cifar10.py` |
+
+## Setup
+
+Same model and hyperparameters everywhere: a VGG-style CNN — three `conv-conv-pool` blocks
+(64→128→256 channels, `3×3` convs + BatchNorm), then `Dense(4096→256) → Dropout(0.5) → Dense(256→10)`
+(~2.2M params) — with `AdamW(1e-3)`, cross-entropy, batch size 128, **10 epochs**, on the **GPU**.
+CIFAR-10 (`uoft-cs/cifar10`) is pulled via HF `datasets`. Augmentation is the classic pipeline:
+per-channel **normalize**, random **crop** to `32×32` with 4-px zero padding, and a random
+**horizontal flip** (test data is only normalized). The Julia side uses **MLUtils 0.4.12**, whose
+`num_workers` path returns collated batches through **shared memory** — see the Distributed rows.
+
+Timings are single-run wall-clock on an **RTX 5090** (CPU: Ryzen Threadripper PRO 9955WX) —
+indicative, not rigorous (expect run-to-run variation). Each script first runs a short verbose
+**DEMO** (accuracy per epoch), then times each config twice: once for **full** training and once for
+**loader iteration alone** (no model, same 10 epochs). **Every timed config discards one warm-up
+epoch before timing**, so Julia's first-call JIT, worker-process startup, and the shm-session build —
+and PyTorch's persistent-worker spawn — stay out of the numbers. These are steady-state per-epoch
+costs. Both frameworks reach ~**80–83%** test accuracy after 10 epochs — this is a data-loading
+timing benchmark, not a numerical match (init/weight-decay details differ), so accuracies only track
+loosely.
+
+In the tables below, **full** = train + load, and **load only** = iterate the `DataLoader` for the
+same 10 epochs consuming each batch but running no model. For the serial paths (no prefetch) `full ≈
+load + compute`, so `full − load` is roughly the GPU step; for the worker/thread paths loading
+overlaps compute, so `full` collapses toward whichever of the two is larger.
+
+## Results
+
+### Julia — Flux + HuggingFaceDatasets.jl (`-t4`, MLUtils 0.4.12)
+
+| config | data loading | full | load only |
+| --- | --- | ---: | ---: |
+| Serial | on-the-fly, `num_workers=0` | 46.9 s | 33.6 s |
+| Serial Materialized | in-memory `[:]`, `num_workers=0` | 16.4 s | 5.5 s |
+| Parallel Materialized | in-memory, `parallel=true` (threads) | 14.5 s | 3.3 s |
+| Distributed | on-the-fly, `num_workers=2` (processes, shared-mem) | 23.7 s | 20.6 s |
+| Distributed | on-the-fly, `num_workers=4` (processes, shared-mem) | 17.4 s | 11.8 s |
+| Distributed | on-the-fly, `num_workers=8` (processes, shared-mem) | 16.4 s | 7.4 s |
+
+### PyTorch (`pytorch_cifar10.py`)
+
+| config | data loading | full | load only |
+| --- | --- | ---: | ---: |
+| Serial | on-the-fly, `num_workers=0` | 53.0 s | 40.1 s |
+| Serial Materialized | in-memory tensors, `num_workers=0` | 16.5 s | 10.1 s |
+| Parallel Materialized | in-memory, `num_workers=4` (processes) | 11.5 s | 3.9 s |
+| Distributed | on-the-fly, `num_workers=2` (processes) | 22.0 s | 22.1 s |
+| Distributed | on-the-fly, `num_workers=4` (processes) | 13.0 s | 12.3 s |
+| Distributed | on-the-fly, `num_workers=8` (processes) | 11.9 s | 7.2 s |
+
+The on-the-fly rows use the standard torchvision order (augment the PIL image, then
+`ToTensor`+`Normalize`); the materialized rows augment already-normalized CHW tensors, matching
+Julia's decode/augment split. The pixels differ slightly, but each is the idiomatic path in its
+framework, so the per-item work is representative.
+
+## Takeaways
+
+- **`num_workers` now scales in Julia — the shared-memory transport (MLUtils 0.4.12).** Julia's
+ Distributed `load only` drops **33.6 → 20.6 → 11.8 → 7.4 s** at `num_workers` 0/2/4/8 (1.6× / 2.8× /
+ 4.5×), tracking PyTorch's own worker curve almost exactly (**40.1 → 22.1 → 12.3 → 7.2 s**). Before
+ 0.4.12 each collated batch came back through a serialize→socket→deserialize round-trip and the curve
+ was flat — more workers didn't help. Now only a shared-memory *handle* crosses the socket (the ~1.5 MB
+ of pixels stays put), so the process-parallel path parallelizes for real. **Julia's distributed loader
+ is now competitive with PyTorch's.**
+- **Julia loads at least as fast as PyTorch, and faster off the distributed path.** Across `load
+ only`: serial 33.6 vs 40.1 s, materialized 5.5 vs 10.1 s, threaded 3.3 vs 3.9 s; distributed is a
+ tie (11.8 vs 12.3; 7.4 vs 7.2). Less per-item Python overhead reaching the Arrow rows.
+- **PyTorch still computes this CNN faster.** Once loading is cheap or fully hidden, the GPU step
+ dominates and PyTorch's mature cuDNN path is quicker: the full-training floor is ~**11.5 s**
+ (Parallel Materialized) vs Julia's ~**14.5 s**. So PyTorch edges the full-training totals on the
+ parallel paths (Distributed 13.0 vs 17.4 s at 4 workers) even though its loading is marginally slower —
+ the two frameworks pull in opposite directions, but the gap is small and it's now the *compute*, not
+ the loading, that decides it.
+- **When the data fits in RAM, materialize + threads is Julia's best path.** Decoding all 50k PNGs
+ once drops per-epoch loading to near-nothing (`load only` 5.5 s materialized; **3.3 s** with
+ `parallel=true`), and `parallel=true` uses **threads** — shared memory, no IPC, no serialization at
+ all — so it is both the fastest loading path here and the fastest Julia full-training config
+ (14.5 s). PyTorch can't thread the loader (GIL), so its in-memory-parallel path uses worker
+ *processes* (3.9 / 11.5 s). Rule of thumb: **materialize + `parallel=true`** when the set fits in
+ memory; reach for **`num_workers`** when it doesn't and the per-epoch decode must be parallelized
+ past the GIL.
diff --git a/perf/flux_cifar10/flux_cifar10.jl b/perf/flux_cifar10/flux_cifar10.jl
new file mode 100644
index 0000000..c41f824
--- /dev/null
+++ b/perf/flux_cifar10/flux_cifar10.jl
@@ -0,0 +1,201 @@
+using Random, Statistics
+using Flux
+using Flux.Losses: logitcrossentropy
+using Flux: onecold, onehotbatch
+using HuggingFaceDatasets
+using MLUtils: MLUtils, mapobs
+using CUDA, cuDNN
+
+# Train on the GPU when one is available (this example is written for it), else fall back to CPU.
+const DEVICE = CUDA.functional() ? gpu : cpu
+
+# CIFAR-10 per-channel mean/std (the standard values), shaped for (W, H, C, N) broadcasting.
+const CIFAR_MEAN = reshape(Float32[0.4914, 0.4822, 0.4465], 1, 1, 3, 1)
+const CIFAR_STD = reshape(Float32[0.2470, 0.2435, 0.2616], 1, 1, 3, 1)
+
+# Decode a raw batch to normalized WHCN Float32. Under the "julia" format the image column is a
+# stacked (C, W, H, N) UInt8 array — channel axis first, from the numpy→Julia axis reversal — so
+# permute to Flux's (W, H, C, N) and standardize per channel. Deterministic, hence materializable.
+function cifar_decode(batch)
+ x = Float32.(batch["img"]) ./ 255f0 # (C, W, H, N)
+ x = permutedims(x, (2, 3, 1, 4)) # (W, H, C, N) — Flux WHCN layout
+ x = (x .- CIFAR_MEAN) ./ CIFAR_STD
+ return (; image = x, label = batch["label"])
+end
+
+# Standard CIFAR-10 training augmentation, applied per batch: zero-pad by 4 and take a random
+# 32×32 crop, then flip horizontally with probability 1/2 — the classic crop+flip pipeline. Pure
+# Julia (no Python), so it parallelizes across threads (`parallel=true`) and worker processes
+# (`num_workers`). It is random per call, so it must run every epoch and must not be materialized.
+function cifar_augment(batch)
+ x = batch.image
+ W, H, C, N = size(x)
+ pad = 4
+ out = similar(x)
+ padded = zeros(Float32, W + 2pad, H + 2pad, C)
+ for n in 1:N
+ fill!(padded, 0f0)
+ @views padded[pad+1:pad+W, pad+1:pad+H, :] .= x[:, :, :, n]
+ i, j = rand(0:2pad), rand(0:2pad) # random top-left crop offset
+ if rand(Bool)
+ @views out[:, :, :, n] .= padded[i+W:-1:i+1, j+1:j+H, :] # crop + horizontal flip
+ else
+ @views out[:, :, :, n] .= padded[i+1:i+W, j+1:j+H, :] # crop
+ end
+ end
+ return (; image = out, label = batch.label)
+end
+
+# On-the-fly training pipeline: decode then augment. A named function (not a closure) so the
+# `num_workers` path can ship it — and the module globals it reads — to worker processes.
+cifar_train_transform(batch) = cifar_augment(cifar_decode(batch))
+
+# A small VGG-style CNN: three (conv-conv-pool) blocks of increasing width, then a classifier.
+# BatchNorm speeds up convergence; both it and Dropout switch behavior between train/eval, which
+# Flux handles automatically (active inside `withgradient`, inactive during plain inference).
+function make_model()
+ return Chain(
+ Conv((3, 3), 3 => 64, pad = 1, bias = false), BatchNorm(64, relu),
+ Conv((3, 3), 64 => 64, pad = 1, bias = false), BatchNorm(64, relu),
+ MaxPool((2, 2)),
+ Conv((3, 3), 64 => 128, pad = 1, bias = false), BatchNorm(128, relu),
+ Conv((3, 3), 128 => 128, pad = 1, bias = false), BatchNorm(128, relu),
+ MaxPool((2, 2)),
+ Conv((3, 3), 128 => 256, pad = 1, bias = false), BatchNorm(256, relu),
+ Conv((3, 3), 256 => 256, pad = 1, bias = false), BatchNorm(256, relu),
+ MaxPool((2, 2)),
+ Flux.flatten,
+ Dense(256 * 4 * 4, 256, relu),
+ Dropout(0.5),
+ Dense(256, 10),
+ )
+end
+
+function loss_and_accuracy(data_loader, model, device)
+ acc = 0
+ ls = 0.0f0
+ num = 0
+ for (x, y) in data_loader
+ x = x |> device
+ ŷ = model(x) |> cpu # logits back on the CPU for cheap metric bookkeeping
+ yoh = onehotbatch(y, 0:9)
+ ls += logitcrossentropy(ŷ, yoh, agg = sum)
+ acc += sum(onecold(ŷ, 0:9) .== y)
+ num += length(y)
+ end
+ return ls / num, acc / num
+end
+
+# `num_workers = 0` loads on the main process; `num_workers > 0` spreads each batch's `getobs`
+# (the CPython image decode) over that many worker processes, sidestepping the GIL — and as of
+# MLUtils 0.4.12 the collated batch returns to the main process through **shared memory** (only a
+# handle crosses the socket, not the ~1.5 MB of pixels), so the process-parallel path now actually
+# scales. `parallel=true` instead uses background threads — useful once the data is materialized,
+# where the per-batch work left is the pure-Julia augmentation (no GIL to serialize it).
+#
+# Returns the wall-clock seconds of `epochs` timed epochs. One extra warm-up epoch runs first and is
+# discarded, so Julia's first-call JIT, worker-process startup, and the shm-session build stay out of
+# the numbers — we measure steady-state per-epoch cost. (The DEMO, `verbose=true`, skips this and
+# just reports accuracy per epoch from initialization.)
+function train(; epochs = 5, num_workers = 0, materialize = false, parallel = false, verbose = true,
+ loader_only = false)
+ batchsize = 128
+ device = DEVICE
+
+ train_ds = load_dataset("uoft-cs/cifar10", split = "train")
+ test_ds = load_dataset("uoft-cs/cifar10", split = "test")
+
+ if materialize
+ # Decode once into memory, then augment per batch on top of the in-memory tensors: the
+ # per-epoch Python decode is gone and only the (pure-Julia) augmentation remains.
+ train_base = mapobs(cifar_decode, train_ds)[:]
+ train_data = mapobs(cifar_augment, train_base)
+ test_data = mapobs(cifar_decode, test_ds)[:]
+ else
+ # Decode + augment on the fly every batch; the CPython decode runs under the GIL.
+ train_data = mapobs(cifar_train_transform, train_ds)
+ test_data = mapobs(cifar_decode, test_ds)
+ end
+
+ train_loader = Flux.DataLoader(train_data; batchsize, shuffle = true, num_workers, parallel)
+ test_loader = Flux.DataLoader(test_data; batchsize, num_workers, parallel)
+
+ if loader_only
+ # Iterate the training loader, consuming each batch but running no model — isolates data
+ # loading (decode + augment + collate + worker IPC) from GPU compute. One warm-up epoch is
+ # discarded, then `epochs` epochs are timed.
+ for (x, y) in train_loader; end # warm-up (discarded)
+ seen = 0
+ return @elapsed for _ in 1:epochs
+ for (x, y) in train_loader
+ seen += length(y)
+ end
+ end
+ end
+
+ model = make_model() |> device
+ opt = Flux.setup(AdamW(1e-3), model)
+
+ function train_epoch!()
+ for (x, y) in train_loader
+ x = x |> device
+ yoh = onehotbatch(y, 0:9) |> device
+ loss, grads = Flux.withgradient(m -> logitcrossentropy(m(x), yoh), model)
+ Flux.update!(opt, model, grads[1])
+ end
+ end
+
+ if verbose
+ # DEMO: report train/test accuracy per epoch from initialization (no timing, no warm-up).
+ function report(epoch)
+ train_loss, train_acc = loss_and_accuracy(train_loader, model, device)
+ test_loss, test_acc = loss_and_accuracy(test_loader, model, device)
+ r(x) = round(x, digits = 3)
+ r(x::Int) = x
+ @info map(r, (; epoch, train_loss, train_acc, test_loss, test_acc))
+ end
+ report(0)
+ for epoch in 1:epochs
+ train_epoch!()
+ report(epoch)
+ end
+ return
+ end
+
+ train_epoch!() # warm-up (discarded)
+ return @elapsed for _ in 1:epochs
+ train_epoch!()
+ end
+end
+
+const EPOCHS = parse(Int, get(ENV, "EPOCHS", "10"))
+
+report_time(t) = println(" ", round(t; digits = 1), " seconds ($EPOCHS epochs, warm-up discarded)")
+
+println("### DEMO — CNN learning CIFAR-10 on ", DEVICE === gpu ? "GPU" : "CPU", " (accuracy per epoch)")
+train(; epochs = EPOCHS, materialize = true, verbose = true)
+
+# Each config below runs one warm-up epoch (discarded) before its timed epochs, so first-call JIT,
+# worker-process startup, and the shm-session build never land in the reported time.
+println("\n#### FULL TRAINING — model + data loading ###########")
+MLUtils.close_dataloader_pool()
+println("### Serial"); report_time(train(; epochs = EPOCHS, num_workers = 0, materialize = false, verbose = false))
+println("### Serial Materialized"); report_time(train(; epochs = EPOCHS, num_workers = 0, materialize = true, verbose = false))
+println("### Parallel Materialized"); report_time(train(; epochs = EPOCHS, num_workers = 0, materialize = true, parallel = true, verbose = false))
+println("### Distributed (2 workers)"); report_time(train(; epochs = EPOCHS, num_workers = 2, materialize = false, verbose = false))
+println("### Distributed (4 workers)"); report_time(train(; epochs = EPOCHS, num_workers = 4, materialize = false, verbose = false))
+println("### Distributed (8 workers)"); report_time(train(; epochs = EPOCHS, num_workers = 8, materialize = false, verbose = false))
+MLUtils.close_dataloader_pool()
+println("#### END FULL TRAINING ###########")
+
+# Same configs, but iterating the loader with no model — the pure data-loading cost. Comparing
+# against the full-training numbers shows how much of each config is loading vs. GPU compute.
+println("\n#### DATA-LOADING ONLY — no model, same $EPOCHS epochs ###########")
+println("### Serial"); report_time(train(; epochs = EPOCHS, num_workers = 0, materialize = false, verbose = false, loader_only = true))
+println("### Serial Materialized"); report_time(train(; epochs = EPOCHS, num_workers = 0, materialize = true, verbose = false, loader_only = true))
+println("### Parallel Materialized"); report_time(train(; epochs = EPOCHS, num_workers = 0, materialize = true, parallel = true, verbose = false, loader_only = true))
+println("### Distributed (2 workers)"); report_time(train(; epochs = EPOCHS, num_workers = 2, materialize = false, verbose = false, loader_only = true))
+println("### Distributed (4 workers)"); report_time(train(; epochs = EPOCHS, num_workers = 4, materialize = false, verbose = false, loader_only = true))
+println("### Distributed (8 workers)"); report_time(train(; epochs = EPOCHS, num_workers = 8, materialize = false, verbose = false, loader_only = true))
+MLUtils.close_dataloader_pool()
+println("#### END DATA-LOADING ONLY ###########")
diff --git a/perf/flux_cifar10/pytorch_cifar10.py b/perf/flux_cifar10/pytorch_cifar10.py
new file mode 100644
index 0000000..29292c3
--- /dev/null
+++ b/perf/flux_cifar10/pytorch_cifar10.py
@@ -0,0 +1,224 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["torch", "torchvision", "datasets", "numpy", "pillow"]
+# ///
+#
+# PyTorch counterpart of flux_cifar10.jl (Flux + HuggingFaceDatasets), for a data-loading timing
+# comparison on the same task: a small VGG-style CNN on CIFAR-10 pulled from the same HuggingFace
+# `datasets` Arrow dataset, trained on the GPU with the standard crop+flip augmentation. Run with:
+# uv run perf/flux_cifar10/pytorch_cifar10.py
+#
+# The configs mirror flux_cifar10.jl. The on-the-fly path uses the idiomatic HF pattern
+# (`Dataset.with_transform` + the standard torchvision crop→flip→ToTensor→Normalize order); the
+# materialized path decodes the split into normalized CHW tensors once and then augments per item.
+# One caveat: PyTorch's only DataLoader parallelism is multiprocess (`num_workers`), so the
+# "Parallel Materialized" row uses worker *processes* over in-memory tensors — Julia's `parallel=true`
+# there uses *threads* (cheap, shared memory), which PyTorch cannot do because of the GIL. This is a
+# timing benchmark, not a numerical match: exact init/weight-decay etc. differ, so accuracies only
+# track loosely.
+
+import functools
+import os
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+from datasets import load_dataset, disable_progress_bars
+from datasets.utils.logging import set_verbosity_error
+from torch.utils.data import DataLoader, Dataset, TensorDataset
+from torchvision import transforms
+
+disable_progress_bars()
+set_verbosity_error()
+
+BATCHSIZE = 128
+EPOCHS = int(os.environ.get("EPOCHS", "10"))
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Standard CIFAR-10 per-channel mean/std (same values as flux_cifar10.jl).
+MEAN = (0.4914, 0.4822, 0.4465)
+STD = (0.2470, 0.2435, 0.2616)
+
+# On-the-fly path — the usual torchvision order: augment the PIL image, then tensor + normalize.
+TRAIN_TF = transforms.Compose([
+ transforms.RandomCrop(32, padding=4), # zero-pad by 4, take a random 32x32 crop
+ transforms.RandomHorizontalFlip(),
+ transforms.ToTensor(), # PIL HWC uint8 -> CHW float in [0, 1]
+ transforms.Normalize(MEAN, STD),
+])
+TEST_TF = transforms.Compose([transforms.ToTensor(), transforms.Normalize(MEAN, STD)])
+
+# Materialized path — tensors are already decoded+normalized CHW floats, so only the spatial
+# augmentation remains (crop pads with zeros in normalized space, matching Julia's materialized path).
+MAT_TRAIN_TF = transforms.Compose([transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()])
+
+
+def make_model():
+ """A small VGG-style CNN: three (conv-conv-pool) blocks, then a classifier. Mirrors flux_cifar10.jl."""
+ def block(ci, co):
+ return [nn.Conv2d(ci, co, 3, padding=1, bias=False), nn.BatchNorm2d(co), nn.ReLU(inplace=True)]
+ return nn.Sequential(
+ *block(3, 64), *block(64, 64), nn.MaxPool2d(2),
+ *block(64, 128), *block(128, 128), nn.MaxPool2d(2),
+ *block(128, 256), *block(256, 256), nn.MaxPool2d(2),
+ nn.Flatten(),
+ nn.Linear(256 * 4 * 4, 256), nn.ReLU(inplace=True),
+ nn.Dropout(0.5),
+ nn.Linear(256, 10),
+ ).to(DEVICE)
+
+
+def apply_transform(batch, tf):
+ """Decode + transform on access (registered with `with_transform`); picklable to workers."""
+ batch["pixel_values"] = [tf(img) for img in batch["img"]]
+ return batch
+
+
+def collate(examples):
+ x = torch.stack([e["pixel_values"] for e in examples])
+ y = torch.tensor([e["label"] for e in examples])
+ return x, y
+
+
+class TensorAugDataset(Dataset):
+ """Materialized path: hold decoded+normalized CHW float tensors in memory, augment per item."""
+
+ def __init__(self, images, labels, transform):
+ self.images, self.labels, self.transform = images, labels, transform
+
+ def __len__(self):
+ return len(self.labels)
+
+ def __getitem__(self, i):
+ return self.transform(self.images[i]), int(self.labels[i])
+
+
+def materialize(hf_ds):
+ """Decode the whole split into in-memory normalized tensors up front (the `[:]` path in Julia)."""
+ ds = hf_ds.with_format("numpy")
+ col = ds["img"]
+ imgs = col if isinstance(col, np.ndarray) else np.stack(col) # (N, H, W, C) uint8
+ imgs = torch.from_numpy(imgs).permute(0, 3, 1, 2).contiguous().float().div_(255.0) # (N, C, H, W)
+ mean = torch.tensor(MEAN).view(1, 3, 1, 1)
+ std = torch.tensor(STD).view(1, 3, 1, 1)
+ imgs = (imgs - mean) / std
+ labels = torch.from_numpy(np.asarray(ds["label"])).long()
+ return imgs, labels
+
+
+@torch.no_grad()
+def loss_and_accuracy(loader, model):
+ model.eval()
+ lossfn = nn.CrossEntropyLoss(reduction="sum")
+ ls, correct, num = 0.0, 0, 0
+ for x, y in loader:
+ x, y = x.to(DEVICE), y.to(DEVICE)
+ logits = model(x)
+ ls += lossfn(logits, y).item()
+ correct += (logits.argmax(1) == y).sum().item()
+ num += y.shape[0]
+ return ls / num, correct / num
+
+
+def make_loaders(num_workers, materialize_data):
+ if materialize_data:
+ train_hf = load_dataset("uoft-cs/cifar10", split="train")
+ test_hf = load_dataset("uoft-cs/cifar10", split="test")
+ train_ds = TensorAugDataset(*materialize(train_hf), MAT_TRAIN_TF)
+ test_ds = TensorDataset(*materialize(test_hf))
+ collate_fn = None
+ else:
+ ds = load_dataset("uoft-cs/cifar10")
+ train_ds = ds["train"].with_transform(functools.partial(apply_transform, tf=TRAIN_TF))
+ test_ds = ds["test"].with_transform(functools.partial(apply_transform, tf=TEST_TF))
+ collate_fn = collate
+
+ # persistent_workers keeps the pool alive across epochs (closer to Julia's leased pool) so the
+ # discarded warm-up epoch below pays the worker-spawn cost once, out of the timed region.
+ kw = dict(num_workers=num_workers, persistent_workers=num_workers > 0,
+ pin_memory=(DEVICE.type == "cuda"), collate_fn=collate_fn)
+ train_loader = DataLoader(train_ds, batch_size=BATCHSIZE, shuffle=True, **kw)
+ test_loader = DataLoader(test_ds, batch_size=BATCHSIZE, **kw)
+ return train_loader, test_loader
+
+
+def train(num_workers=0, materialize_data=False, verbose=False, loader_only=False):
+ """Return wall-clock seconds of EPOCHS timed epochs; one warm-up epoch runs first and is discarded
+ (spawns the persistent workers, warms cuDNN), so worker-spawn cost stays out of the timing.
+ `verbose=True` is the DEMO path: report accuracy per epoch instead, no warm-up, no timing."""
+ train_loader, test_loader = make_loaders(num_workers, materialize_data)
+
+ if loader_only:
+ for _ in train_loader: # warm-up (discarded)
+ pass
+ seen = 0
+ t0 = time.perf_counter()
+ for _ in range(EPOCHS):
+ for x, y in train_loader:
+ seen += y.shape[0]
+ return time.perf_counter() - t0
+
+ model = make_model()
+ opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
+ lossfn = nn.CrossEntropyLoss()
+
+ def run_epoch():
+ model.train()
+ for x, y in train_loader:
+ x, y = x.to(DEVICE), y.to(DEVICE)
+ opt.zero_grad()
+ loss = lossfn(model(x), y)
+ loss.backward()
+ opt.step()
+
+ if verbose:
+ def report(epoch):
+ tr_loss, tr_acc = loss_and_accuracy(train_loader, model)
+ te_loss, te_acc = loss_and_accuracy(test_loader, model)
+ print(f"(epoch = {epoch}, train_loss = {tr_loss:.3f}, train_acc = {tr_acc:.3f}, "
+ f"test_loss = {te_loss:.3f}, test_acc = {te_acc:.3f})")
+ report(0)
+ for epoch in range(1, EPOCHS + 1):
+ run_epoch()
+ report(epoch)
+ return
+
+ run_epoch() # warm-up (discarded)
+ t0 = time.perf_counter()
+ for _ in range(EPOCHS):
+ run_epoch()
+ return time.perf_counter() - t0
+
+
+def timed(name, **kw):
+ t = train(**kw)
+ print(f"### {name}\n {t:.1f} seconds ({EPOCHS} epochs, warm-up discarded)")
+
+
+if __name__ == "__main__":
+ dev = "GPU" if DEVICE.type == "cuda" else "CPU"
+ print(f"### DEMO — CNN learning CIFAR-10 on {dev} (accuracy per epoch)")
+ train(num_workers=0, materialize_data=True, verbose=True)
+
+ # Each config runs one warm-up epoch (discarded) before its timed epochs, so worker-spawn and
+ # cuDNN autotune stay out of the numbers — matching the Julia script's warm-up discipline.
+ print("\n#### FULL TRAINING — model + data loading ###########")
+ timed("Serial", num_workers=0, materialize_data=False)
+ timed("Serial Materialized", num_workers=0, materialize_data=True)
+ timed("Parallel Materialized", num_workers=4, materialize_data=True)
+ timed("Distributed (2 workers)", num_workers=2, materialize_data=False)
+ timed("Distributed (4 workers)", num_workers=4, materialize_data=False)
+ timed("Distributed (8 workers)", num_workers=8, materialize_data=False)
+ print("#### END FULL TRAINING ###########")
+
+ # Same configs, but iterating the loader with no model — the pure data-loading cost. Comparing
+ # against the full-training numbers shows how much of each config is loading vs. GPU compute.
+ print(f"\n#### DATA-LOADING ONLY — no model, same {EPOCHS} epochs ###########")
+ timed("Serial", num_workers=0, materialize_data=False, loader_only=True)
+ timed("Serial Materialized", num_workers=0, materialize_data=True, loader_only=True)
+ timed("Parallel Materialized", num_workers=4, materialize_data=True, loader_only=True)
+ timed("Distributed (2 workers)", num_workers=2, materialize_data=False, loader_only=True)
+ timed("Distributed (4 workers)", num_workers=4, materialize_data=False, loader_only=True)
+ timed("Distributed (8 workers)", num_workers=8, materialize_data=False, loader_only=True)
+ print("#### END DATA-LOADING ONLY ###########")
diff --git a/examples/flux_mnist/Project.toml b/perf/flux_mnist/Project.toml
similarity index 96%
rename from examples/flux_mnist/Project.toml
rename to perf/flux_mnist/Project.toml
index eeed48e..a33aec1 100644
--- a/examples/flux_mnist/Project.toml
+++ b/perf/flux_mnist/Project.toml
@@ -13,4 +13,4 @@ HuggingFaceDatasets = {path = "../.."}
[compat]
DataStructures = "0.19"
-MLUtils = "0.4"
\ No newline at end of file
+MLUtils = "0.4.12"
\ No newline at end of file
diff --git a/perf/flux_mnist/README.md b/perf/flux_mnist/README.md
new file mode 100644
index 0000000..967d95c
--- /dev/null
+++ b/perf/flux_mnist/README.md
@@ -0,0 +1,88 @@
+# Flux MNIST — data-loading benchmark (CPU)
+
+Trains a small MLP on MNIST pulled from the HuggingFace `datasets` library, comparing
+**data-loading strategies** (on-the-fly vs. materialized, thread- vs. process-parallel) in
+Flux+HuggingFaceDatasets.jl against an equivalent PyTorch version. It's the tiny-model counterpart of
+the [`flux_cifar10`](../flux_cifar10) example: same structure, but the MLP is so small that this
+benchmark is almost purely about **data loading**, not compute.
+
+## Files
+
+| file | stack | notes |
+| --- | --- | --- |
+| [`flux_mnist.jl`](flux_mnist.jl) | Flux + HuggingFaceDatasets.jl | `julia --project=. -t4 flux_mnist.jl` |
+| [`pytorch_mnist.py`](pytorch_mnist.py) | PyTorch + `datasets` | `uv run pytorch_mnist.py` |
+
+## Setup
+
+Same model and hyperparameters everywhere: MLP `784 → 100 → 100 → 10`, `AdamW(1e-3)`, cross-entropy,
+batch size 128, **10 epochs**, **CPU only** (the model is tiny — keeping it off the GPU isolates the
+data-loading cost, which is the point of this example). MNIST (`ylecun/mnist`) is pulled via HF
+`datasets`; the Julia side uses **MLUtils 0.4.12**, whose `num_workers` path returns collated batches
+through **shared memory**. Both frameworks reach ~**98%** test accuracy after 10 epochs.
+
+Timings are single-run wall-clock on a **Ryzen Threadripper PRO 9955WX** (Julia `-t4`) — indicative,
+not rigorous (expect run-to-run variation). Each script first runs a short verbose **DEMO** (accuracy
+per epoch), then times each config twice: once for **full** training and once for **loader iteration
+alone** (no model, same 10 epochs). **Every timed config discards one warm-up epoch before timing**, so
+Julia's first-call JIT, worker-process startup, and the shm-session build — and PyTorch's
+persistent-worker spawn — stay out of the numbers. These are steady-state per-epoch costs.
+
+> **Threads + a PythonCall-backed dataset.** HuggingFaceDatasets shares numpy arrays into Julia
+> *zero-copy*, keeping the batches backed by Python-owned buffers. Those buffers are released through
+> PythonCall's GIL-deferred decref: a finalizer that can't take the GIL just enqueues the pointer, so
+> one firing on a `parallel=true` worker thread never deadlocks. The materialized paths here decode the
+> split into plain in-memory Julia arrays with `[:]` (no Python buffer left to touch), which is what
+> lets `parallel=true` (threads) run cleanly on them — see the comments in
+> [`flux_mnist.jl`](flux_mnist.jl). For on-the-fly PythonCall data, prefer `num_workers` (separate
+> processes, separate GILs) over `parallel` (threads) to get real parallelism past the GIL.
+
+## Results
+
+### Julia — Flux + HuggingFaceDatasets.jl (`-t4`, MLUtils 0.4.12)
+
+| config | data loading | full | load only |
+| --- | --- | ---: | ---: |
+| Serial | on-the-fly, `num_workers=0` | 23.0 s | 17.8 s |
+| Serial Materialized | in-memory `[:]`, `num_workers=0` | 4.7 s | 0.4 s |
+| Parallel Materialized | in-memory, `parallel=true` (threads) | 5.0 s | 0.3 s |
+| Distributed | on-the-fly, `num_workers=2` (processes, shared-mem) | 13.3 s | 10.0 s |
+| Distributed | on-the-fly, `num_workers=4` (processes, shared-mem) | 7.9 s | 5.4 s |
+| Distributed | on-the-fly, `num_workers=8` (processes, shared-mem) | 5.5 s | 3.1 s |
+
+### PyTorch (`pytorch_mnist.py`)
+
+| config | data loading | full | load only |
+| --- | --- | ---: | ---: |
+| Serial | on-the-fly, `num_workers=0` | 17.9 s | 13.4 s |
+| Serial Materialized | in-memory tensors, `num_workers=0` | 4.4 s | 1.3 s |
+| Parallel Materialized | in-memory, `num_workers=4` (processes) | 5.2 s | 1.2 s |
+| Distributed | on-the-fly, `num_workers=2` (processes) | 9.8 s | 8.3 s |
+| Distributed | on-the-fly, `num_workers=4` (processes) | 6.2 s | 4.9 s |
+| Distributed | on-the-fly, `num_workers=8` (processes) | 6.4 s | 2.8 s |
+
+## Takeaways
+
+- **For data this small, materialize and you're done.** Decoding all 60k images once drops per-epoch
+ loading to essentially nothing — `load only` **17.8 → 0.4 s** in Julia, **13.4 → 1.3 s** in PyTorch —
+ because MNIST is ~180 MB and lives comfortably in RAM. Materialized is the fastest full-training
+ config in both frameworks (Julia 4.7 s, PyTorch 4.4 s).
+- **`num_workers` now scales — it's just unnecessary here.** With the shared-memory transport, Julia's
+ Distributed `load only` drops **17.8 → 10.0 → 5.4 → 3.1 s** at `num_workers` 0/2/4/8, tracking PyTorch
+ (13.4 → 8.3 → 4.9 → 2.8 s) almost step for step. This is a real change from the past, when process
+ workers were *slower* than serial for a model this small. But worker IPC/startup still can't beat a
+ materialized in-memory scan (0.3–1.3 s), so for a dataset that fits in memory, **materialize beats
+ `num_workers`** — reach for workers only when the data *doesn't* fit and the per-epoch decode must be
+ parallelized.
+- **The MLP is too small for compute to matter — this is a loading benchmark.** `full − load` is ~3–5 s
+ everywhere (the entire train step for a 784→100→100→10 MLP is nearly free), so the two frameworks land
+ close together across the board. Julia's materialized iteration is faster (`load only` 0.4 vs 1.3 s);
+ PyTorch's on-the-fly serial is faster (13.4 vs 17.8 s). Effectively a wash — as expected when there's
+ almost nothing to compute. For the case where compute *does* matter, see the
+ [CIFAR-10 example](../flux_cifar10).
+- **`parallel=true` (threads) needs a pure-Julia dataset.** It shines on the materialized path
+ (`load only` 0.3 s, edging serial-materialized — there's nothing left to parallelize once the data
+ is in memory) but only because the materialize step copies the data out of Python first. Threads over
+ a live PythonCall-backed dataset would still be GIL-serialized on every decode (though, thanks to the
+ deferred decref, they no longer deadlock — see the note above); that's what `num_workers` (processes)
+ is for.
diff --git a/perf/flux_mnist/flux_mnist.jl b/perf/flux_mnist/flux_mnist.jl
new file mode 100644
index 0000000..ee7d4b3
--- /dev/null
+++ b/perf/flux_mnist/flux_mnist.jl
@@ -0,0 +1,147 @@
+using Random, Statistics
+using Flux
+using Flux.Losses: logitcrossentropy
+using Flux: onecold, onehotbatch
+using HuggingFaceDatasets
+using MLUtils: MLUtils, mapobs
+
+# The image column is a stacked (W, H, N) UInt8 array, so just rescale to Float32 in [0, 1].
+function mnist_transform(batch)
+ image = batch["image"] ./ 255f0
+ return (; image, label = batch["label"])
+end
+
+function loss_and_accuracy(data_loader, model, device)
+ acc = 0
+ ls = 0.0f0
+ num = 0
+ for (x, y) in data_loader
+ x = x |> device
+ yoh = onehotbatch(y, 0:9) |> device
+ ŷ = model(x)
+ ls += logitcrossentropy(ŷ, yoh, agg = sum)
+ acc += sum(onecold(ŷ, 0:9) .== y)
+ num += length(y)
+ end
+ return ls / num, acc / num
+end
+
+# `num_workers = 0` loads on the main process; `num_workers > 0` spreads each batch's `getobs`
+# (and the CPython read it triggers) over that many worker processes, sidestepping the GIL — and as
+# of MLUtils 0.4.12 the collated batch returns to the main process through **shared memory** (only a
+# handle crosses the socket, not the pixels). `parallel=true` instead uses background threads. This
+# is a tiny MLP, though, so the per-batch worker/thread overhead tends to outweigh any parallel-load
+# gain — see the README.
+#
+# Returns the wall-clock seconds of `epochs` timed epochs. One extra warm-up epoch runs first and is
+# discarded, so Julia's first-call JIT, worker-process startup, and the shm-session build stay out of
+# the numbers — we measure steady-state per-epoch cost. (The DEMO, `verbose=true`, skips this and
+# just reports accuracy per epoch.)
+function train(; epochs = 4, num_workers = 0, materialize = false, parallel = false, verbose = true,
+ loader_only = false)
+ batchsize = 128
+ nhidden = 100
+ device = cpu
+
+ train_data = load_dataset("ylecun/mnist", split = "train")
+ test_data = load_dataset("ylecun/mnist", split = "test")
+ # Apply the transform lazily so it runs per batch during iteration (on the workers when
+ # `num_workers > 0`); `mapobs`/`ObsView`-wrapped datasets compose with `num_workers`.
+ train_data = mapobs(mnist_transform, train_data)
+ test_data = mapobs(mnist_transform, test_data)
+ if materialize
+ train_data = train_data[:]
+ test_data = test_data[:]
+ end
+
+ train_loader = Flux.DataLoader(train_data; batchsize, shuffle = true, num_workers, parallel)
+ test_loader = Flux.DataLoader(test_data; batchsize, num_workers, parallel)
+
+ if loader_only
+ # Iterate the training loader, consuming each batch but running no model — isolates data
+ # loading (read + transform + collate + worker IPC) from compute. One warm-up epoch, then
+ # time `epochs` epochs.
+ for (x, y) in train_loader; end # warm-up (discarded)
+ seen = 0
+ return @elapsed for _ in 1:epochs
+ for (x, y) in train_loader
+ seen += length(y)
+ end
+ end
+ end
+
+ model = Chain(Flux.flatten,
+ Dense(28 * 28, nhidden, relu),
+ Dense(nhidden, nhidden, relu),
+ Dense(nhidden, 10)) |> device
+ opt = Flux.setup(AdamW(1e-3), model)
+
+ function train_epoch!()
+ for (x, y) in train_loader
+ x = x |> device
+ yoh = onehotbatch(y, 0:9) |> device
+ loss, grads = Flux.withgradient(m -> logitcrossentropy(m(x), yoh), model)
+ Flux.update!(opt, model, grads[1])
+ end
+ end
+
+ if verbose
+ # DEMO: report train/test accuracy per epoch from initialization (no timing, no warm-up).
+ function report(epoch)
+ train_loss, train_acc = loss_and_accuracy(train_loader, model, device)
+ test_loss, test_acc = loss_and_accuracy(test_loader, model, device)
+ r(x) = round(x, digits = 3)
+ r(x::Int) = x
+ @info map(r, (; epoch, train_loss, train_acc, test_loss, test_acc))
+ end
+ report(0)
+ for epoch in 1:epochs
+ train_epoch!()
+ report(epoch)
+ end
+ return
+ end
+
+ train_epoch!() # warm-up (discarded)
+ return @elapsed for _ in 1:epochs
+ train_epoch!()
+ end
+end
+
+const EPOCHS = parse(Int, get(ENV, "EPOCHS", "4"))
+
+# Print the config name (flushed) before running it and the time after, so progress is visible live
+# even when stdout is redirected to a file — Julia block-buffers a non-TTY stdout otherwise.
+function timed(name; kwargs...)
+ print("### ", rpad(name, 26)); flush(stdout)
+ t = train(; epochs = EPOCHS, verbose = false, kwargs...)
+ println(round(t; digits = 1), " seconds ($EPOCHS epochs, warm-up discarded)"); flush(stdout)
+end
+
+println("### DEMO — MLP learning MNIST on CPU (accuracy per epoch)"); flush(stdout)
+train(; epochs = EPOCHS, materialize = true, verbose = true)
+
+# Each config below runs one warm-up epoch (discarded) before its timed epochs, so first-call JIT,
+# worker-process startup, and the shm-session build never land in the reported time.
+println("\n#### FULL TRAINING — model + data loading ###########"); flush(stdout)
+MLUtils.close_dataloader_pool()
+timed("Serial"; num_workers = 0, materialize = false)
+timed("Serial Materialized"; num_workers = 0, materialize = true)
+timed("Parallel Materialized"; num_workers = 0, materialize = true, parallel = true)
+timed("Distributed (2 workers)"; num_workers = 2, materialize = false)
+timed("Distributed (4 workers)"; num_workers = 4, materialize = false)
+timed("Distributed (8 workers)"; num_workers = 8, materialize = false)
+MLUtils.close_dataloader_pool()
+println("#### END FULL TRAINING ###########"); flush(stdout)
+
+# Same configs, but iterating the loader with no model — the pure data-loading cost. Comparing
+# against the full-training numbers shows how much of each config is loading vs. compute.
+println("\n#### DATA-LOADING ONLY — no model, same $EPOCHS epochs ###########"); flush(stdout)
+timed("Serial"; num_workers = 0, materialize = false, loader_only = true)
+timed("Serial Materialized"; num_workers = 0, materialize = true, loader_only = true)
+timed("Parallel Materialized"; num_workers = 0, materialize = true, parallel = true, loader_only = true)
+timed("Distributed (2 workers)"; num_workers = 2, materialize = false, loader_only = true)
+timed("Distributed (4 workers)"; num_workers = 4, materialize = false, loader_only = true)
+timed("Distributed (8 workers)"; num_workers = 8, materialize = false, loader_only = true)
+MLUtils.close_dataloader_pool()
+println("#### END DATA-LOADING ONLY ###########"); flush(stdout)
diff --git a/perf/flux_mnist/pytorch_mnist.py b/perf/flux_mnist/pytorch_mnist.py
new file mode 100644
index 0000000..aafbd9d
--- /dev/null
+++ b/perf/flux_mnist/pytorch_mnist.py
@@ -0,0 +1,178 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["torch", "datasets", "numpy", "pillow"]
+# ///
+#
+# PyTorch counterpart of flux_mnist.jl (Flux + HuggingFaceDatasets), for a data-loading timing
+# comparison on the same task: an MLP on MNIST pulled from the same HuggingFace `datasets` Arrow
+# dataset, on the CPU. Run with: uv run perf/flux_mnist/pytorch_mnist.py
+#
+# The configs mirror flux_mnist.jl. The on-the-fly path uses the idiomatic HF pattern
+# (`Dataset.with_transform` + a `collate_fn`); the materialized path decodes the split into in-memory
+# tensors once. One caveat: PyTorch's only DataLoader parallelism is multiprocess (`num_workers`), so
+# the "Parallel Materialized" row uses worker *processes* over in-memory tensors — Julia's
+# `parallel=true` there uses *threads* (cheap, shared memory), which PyTorch cannot do because of the
+# GIL. This is a timing benchmark, not a numerical match: the exact AdamW weight decay etc. differ,
+# so accuracies only track loosely.
+
+import os
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+from datasets import load_dataset, disable_progress_bars
+from datasets.utils.logging import set_verbosity_error
+from torch.utils.data import DataLoader, TensorDataset
+
+disable_progress_bars()
+set_verbosity_error()
+
+BATCHSIZE = 128
+NHIDDEN = 100
+EPOCHS = int(os.environ.get("EPOCHS", "4"))
+DEVICE = torch.device("cpu")
+
+
+def make_model():
+ return nn.Sequential(
+ nn.Flatten(),
+ nn.Linear(28 * 28, NHIDDEN), nn.ReLU(),
+ nn.Linear(NHIDDEN, NHIDDEN), nn.ReLU(),
+ nn.Linear(NHIDDEN, 10),
+ ).to(DEVICE)
+
+
+def lazy_transform(batch):
+ """Decode + rescale on access (registered with `with_transform`); picklable to workers.
+ Matches Julia's `image ./ 255f0` (plain [0, 1] scaling, no normalization)."""
+ batch["pixel_values"] = [np.asarray(img, dtype=np.float32) / 255.0 for img in batch["image"]]
+ return batch
+
+
+def collate(examples):
+ x = torch.from_numpy(np.stack([e["pixel_values"] for e in examples]))
+ y = torch.tensor([e["label"] for e in examples])
+ return x, y
+
+
+def materialize(hf_ds):
+ """Decode the whole split into in-memory tensors up front (the `[:]` path in flux_mnist.jl)."""
+ ds = hf_ds.with_format("numpy")
+ col = ds["image"]
+ images = col if isinstance(col, np.ndarray) else np.stack(col)
+ images = torch.from_numpy(images.astype(np.float32) / 255.0) # (N, 28, 28)
+ labels = torch.from_numpy(np.asarray(ds["label"]).astype(np.int64)) # (N,)
+ return TensorDataset(images, labels)
+
+
+def make_loaders(num_workers, materialize_data):
+ if materialize_data:
+ train_ds = materialize(load_dataset("ylecun/mnist", split="train"))
+ test_ds = materialize(load_dataset("ylecun/mnist", split="test"))
+ collate_fn = None
+ else:
+ ds = load_dataset("ylecun/mnist")
+ train_ds = ds["train"].with_transform(lazy_transform)
+ test_ds = ds["test"].with_transform(lazy_transform)
+ collate_fn = collate
+
+ # persistent_workers keeps the pool alive across epochs (closer to Julia's leased pool) so the
+ # discarded warm-up epoch below pays the worker-spawn cost once, out of the timed region.
+ kw = dict(num_workers=num_workers, persistent_workers=num_workers > 0, collate_fn=collate_fn)
+ train_loader = DataLoader(train_ds, batch_size=BATCHSIZE, shuffle=True, **kw)
+ test_loader = DataLoader(test_ds, batch_size=BATCHSIZE, **kw)
+ return train_loader, test_loader
+
+
+@torch.no_grad()
+def loss_and_accuracy(loader, model):
+ model.eval()
+ lossfn = nn.CrossEntropyLoss(reduction="sum")
+ ls, correct, num = 0.0, 0, 0
+ for x, y in loader:
+ x, y = x.to(DEVICE), y.to(DEVICE)
+ logits = model(x)
+ ls += lossfn(logits, y).item()
+ correct += (logits.argmax(1) == y).sum().item()
+ num += y.shape[0]
+ return ls / num, correct / num
+
+
+def train(num_workers=0, materialize_data=False, verbose=False, loader_only=False):
+ """Return wall-clock seconds of EPOCHS timed epochs; one warm-up epoch runs first and is discarded
+ (spawns the persistent workers), so worker-spawn cost stays out of the timing. `verbose=True` is
+ the DEMO path: report accuracy per epoch instead, no warm-up, no timing."""
+ train_loader, test_loader = make_loaders(num_workers, materialize_data)
+
+ if loader_only:
+ for _ in train_loader: # warm-up (discarded)
+ pass
+ seen = 0
+ t0 = time.perf_counter()
+ for _ in range(EPOCHS):
+ for x, y in train_loader:
+ seen += y.shape[0]
+ return time.perf_counter() - t0
+
+ model = make_model()
+ opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
+ lossfn = nn.CrossEntropyLoss()
+
+ def run_epoch():
+ model.train()
+ for x, y in train_loader:
+ x, y = x.to(DEVICE), y.to(DEVICE)
+ opt.zero_grad()
+ loss = lossfn(model(x), y)
+ loss.backward()
+ opt.step()
+
+ if verbose:
+ def report(epoch):
+ tr_loss, tr_acc = loss_and_accuracy(train_loader, model)
+ te_loss, te_acc = loss_and_accuracy(test_loader, model)
+ print(f"(epoch = {epoch}, train_loss = {tr_loss:.3f}, train_acc = {tr_acc:.3f}, "
+ f"test_loss = {te_loss:.3f}, test_acc = {te_acc:.3f})")
+ report(0)
+ for epoch in range(1, EPOCHS + 1):
+ run_epoch()
+ report(epoch)
+ return
+
+ run_epoch() # warm-up (discarded)
+ t0 = time.perf_counter()
+ for _ in range(EPOCHS):
+ run_epoch()
+ return time.perf_counter() - t0
+
+
+def timed(name, **kw):
+ t = train(**kw)
+ print(f"### {name}\n {t:.1f} seconds ({EPOCHS} epochs, warm-up discarded)")
+
+
+if __name__ == "__main__":
+ print("### DEMO — MLP learning MNIST on CPU (accuracy per epoch)")
+ train(num_workers=0, materialize_data=True, verbose=True)
+
+ # Each config runs one warm-up epoch (discarded) before its timed epochs, so worker-spawn stays
+ # out of the numbers — matching the Julia script's warm-up discipline.
+ print("\n#### FULL TRAINING — model + data loading ###########")
+ timed("Serial", num_workers=0, materialize_data=False)
+ timed("Serial Materialized", num_workers=0, materialize_data=True)
+ timed("Parallel Materialized", num_workers=4, materialize_data=True)
+ timed("Distributed (2 workers)", num_workers=2, materialize_data=False)
+ timed("Distributed (4 workers)", num_workers=4, materialize_data=False)
+ timed("Distributed (8 workers)", num_workers=8, materialize_data=False)
+ print("#### END FULL TRAINING ###########")
+
+ # Same configs, but iterating the loader with no model — the pure data-loading cost.
+ print(f"\n#### DATA-LOADING ONLY — no model, same {EPOCHS} epochs ###########")
+ timed("Serial", num_workers=0, materialize_data=False, loader_only=True)
+ timed("Serial Materialized", num_workers=0, materialize_data=True, loader_only=True)
+ timed("Parallel Materialized", num_workers=4, materialize_data=True, loader_only=True)
+ timed("Distributed (2 workers)", num_workers=2, materialize_data=False, loader_only=True)
+ timed("Distributed (4 workers)", num_workers=4, materialize_data=False, loader_only=True)
+ timed("Distributed (8 workers)", num_workers=8, materialize_data=False, loader_only=True)
+ print("#### END DATA-LOADING ONLY ###########")
diff --git a/perf/Project.toml b/perf/load_dataset/Project.toml
similarity index 100%
rename from perf/Project.toml
rename to perf/load_dataset/Project.toml
diff --git a/perf/README.md b/perf/load_dataset/README.md
similarity index 96%
rename from perf/README.md
rename to perf/load_dataset/README.md
index 7816c6e..f390c9c 100644
--- a/perf/README.md
+++ b/perf/load_dataset/README.md
@@ -97,14 +97,14 @@ copy), so the reads run concurrently. Single run; the `serial` row is the in-run
- **The wrapper's overhead is small at batch granularity.** For the realistic `epoch`
pipeline the `julia` format costs ~370 ms vs ~310 ms for Python's `numpy` format — about
- **1.2×**. The convenience of native Julia arrays (copyless `py2jl` via DLPack) is nearly
+ **1.2×**. The convenience of native Julia arrays (copyless, zero-copy `py2jl`) is nearly
free once you work in batches. Per-sample access is where the FFI cost shows up
(`single`: ~2× Python), so avoid it in hot loops.
- **Arrow-backed access is the bottleneck, not Julia.** `MLDatasets` — native, fully
in-memory — is ~100× faster than any `datasets`-backed path. When a dataset fits in
memory, materialize it once with `getobs(ds, :)` and feed an in-memory `DataLoader`
- (see [`../examples/flux_mnist.jl`](../examples/flux_mnist.jl)); the Arrow decode then
+ (see [`../flux_mnist/flux_mnist.jl`](../flux_mnist/flux_mnist.jl)); the Arrow decode then
becomes a one-time cost rather than a per-batch one.
- **When it doesn't fit in memory, parallelize the read with processes, not threads.** The
diff --git a/perf/perf.jl b/perf/load_dataset/perf.jl
similarity index 100%
rename from perf/perf.jl
rename to perf/load_dataset/perf.jl
diff --git a/perf/perf.py b/perf/load_dataset/perf.py
similarity index 100%
rename from perf/perf.py
rename to perf/load_dataset/perf.py
diff --git a/src/HuggingFaceDatasets.jl b/src/HuggingFaceDatasets.jl
index 9fbfc42..be1cc73 100644
--- a/src/HuggingFaceDatasets.jl
+++ b/src/HuggingFaceDatasets.jl
@@ -3,7 +3,6 @@ module HuggingFaceDatasets
using PythonCall
using Compat: @compat
using MLCore: getobs
-using DLPack: DLPack
using ImageCore: colorview, RGB, Gray, N0f8
using Tables: Tables
diff --git a/src/dataset.jl b/src/dataset.jl
index 2eff495..1f5e718 100644
--- a/src/dataset.jl
+++ b/src/dataset.jl
@@ -428,7 +428,7 @@ removes all formatting (raw Python observations); any other string is forwarded
function set_format!(ds::Dataset, format)
if format == "julia"
# Use the numpy format so numeric array columns decode as real N-D arrays and a
- # range index stacks rows into an `(N, dims…)` tensor; `py2jl` (via DLPack) then
+ # range index stacks rows into an `(N, dims…)` tensor; `py2jl` (zero-copy) then
# reverses the axes to a Julia `(dims…, N)` array with the observation axis last.
ds.py.set_format("numpy")
ds.jltransform = py2jl
diff --git a/src/transforms.jl b/src/transforms.jl
index 0460be7..5fb1bf4 100644
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -25,9 +25,9 @@ julia> py2jl(pytuple((1, pylist([2, 3]))))
"""
py2jl(x) = pyconvert(Any, x)
-# Whether a numpy array's dtype is one DLPack (and hence `numpy2jl`) can share: bool,
+# Whether a numpy array's dtype is one `numpy2jl` can share zero-copy: bool,
# signed/unsigned integer, float, or complex. Strings/objects/datetimes are excluded.
-_is_dlpack_numeric(x::Py) = pyconvert(String, x.dtype.kind) in ("b", "i", "u", "f", "c")
+_is_numeric_dtype(x::Py) = pyconvert(String, x.dtype.kind) in ("b", "i", "u", "f", "c")
function py2jl(x::Py)
# handle datasets
@@ -61,10 +61,10 @@ function py2jl(x::Py)
# branch below, so a scalar reads back as a scalar rather than a `fill(x)` 0-d array.
if pyconvert(Int, x.ndim) == 0
return py2jl(x.item())
- # DLPack (`numpy2jl`) only supports numeric dtypes. Non-numeric arrays (strings,
+ # Zero-copy sharing (`numpy2jl`) only supports numeric dtypes. Non-numeric arrays (strings,
# `object` arrays from ragged columns, datetimes, ...) fall back to a nested-list
# conversion, so a string column still comes back as a `Vector{String}`.
- elseif _is_dlpack_numeric(x)
+ elseif _is_numeric_dtype(x)
return numpy2jl(x)
else
return py2jl(x.tolist())
@@ -91,12 +91,25 @@ function py2jl(x::Py)
end
+# Roots the Python buffer backing each zero-copy `numpy2jl` array for exactly as long as the
+# Julia `Array` that views it. The wrapper `Array` is the (weak) key, so an entry — and the
+# Python reference it holds — is dropped automatically once that array is garbage-collected;
+# `WeakKeyDict` is internally locked, so concurrent inserts/evictions are thread-safe.
+#
+# Cleanup routes through PythonCall's GIL-deferred decref (a finalizer that can't take the GIL
+# just enqueues the pointer), so — unlike DLPack's finalizer, which eagerly re-acquires the GIL
+# — a buffer freed on a `DataLoader` worker thread can never deadlock against a thread compiling
+# under the GIL.
+const _NUMPY_BUFFERS = WeakKeyDict{AbstractArray,Any}()
+
"""
numpy2jl(x)
-Convert a numpy array to a Julia array using DLPack.jl.
-The conversion is copyless, and mutations to the Julia array are reflected in the numpy array.
-For row major python arrays, the returned Julia array has permuted dimensions.
+Convert a numpy array to a Julia `Array` sharing memory zero-copy. Mutations to the Julia array
+are reflected in the numpy array (and vice versa). Since numpy is row-major and Julia is
+column-major, the returned array has permuted (reversed) dimensions.
+
+Read-only or non-contiguous numpy buffers cannot be shared safely and are copied first.
This function is called by [`py2jl`](@ref).
See also [`jl2numpy`](@ref).
@@ -113,14 +126,23 @@ julia> numpy2jl(y) # back to a 2×3 Julia array
```
"""
function numpy2jl(x::Py)
- # DLPack cannot import a read-only numpy buffer (numpy >= 2.1 signals read-only, which
- # this DLPack version does not support), and the numpy format hands back read-only
- # arrays for some columns. Copy to a writable array first in that case; the copy is
- # C-contiguous, so `from_dlpack` still reverses the axes and the orientation is correct.
- if !pyconvert(Bool, x.flags.writeable)
+ # `unsafe_wrap` below reinterprets the buffer as a column-major contiguous `Array`, which is
+ # only valid when `x.T` is F-contiguous, i.e. when `x` is C-contiguous. Copy to a fresh,
+ # writable C-contiguous array otherwise. This also covers read-only buffers (numpy >= 2.1
+ # marks some columns read-only): aliasing read-only memory with a writable Julia array would
+ # be unsound. `x.copy()` defaults to C order, so `x.T` is F-contiguous afterwards.
+ if !pyconvert(Bool, x.flags.c_contiguous) || !pyconvert(Bool, x.flags.writeable)
x = x.copy()
end
- return DLPack.from_dlpack(x)
+ # `PyArray(x.T)` is a zero-copy view with reversed axes; it holds the Python reference that
+ # keeps the buffer alive. `unsafe_wrap` then exposes that same memory as a genuine `Array`, so
+ # the result stays on Julia's fast paths (BLAS, GPU host→device copies, linear indexing) that
+ # a `PyArray` — not being a `DenseArray` — would miss. `_NUMPY_BUFFERS` roots the `PyArray` for
+ # the wrapper's lifetime so the buffer outlives every view of it.
+ p = PyArray(x.T)
+ arr = unsafe_wrap(Array, pointer(p), size(p); own = false)
+ _NUMPY_BUFFERS[arr] = p
+ return arr
end
"""