From 5d91b50a3cb61acd6a0aeb6fd1c26c4a16863b38 Mon Sep 17 00:00:00 2001 From: CarloLucibello Date: Tue, 7 Jul 2026 11:59:52 +0200 Subject: [PATCH] Drop DLPack dependency; reorganize examples as perf benchmarks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the DLPack.jl-based numpy→Julia conversion in `numpy2jl` with a vendored zero-copy `PyArray` + `unsafe_wrap`, rooted through a `WeakKeyDict` for lifetime safety. Behavior and performance match DLPack (a genuine `DenseArray`, so BLAS/GPU fast paths are preserved), and as a bonus the GIL-finalizer deadlock is gone: PythonCall's deferred decref means a finalizer firing on a worker thread enqueues the pointer instead of re-acquiring the GIL. Removes the now-unmaintained DLPack dependency. Move the Flux MNIST/CIFAR-10 examples under `perf/` alongside `load_dataset/`, add a 2-worker Distributed row, run the loaders at `-t4`, and refresh both READMEs with fresh 10-epoch benchmark numbers (Threadripper CPU / RTX 5090 GPU). Update all example path references across the docs and READMEs. Co-Authored-By: Claude Opus 4.8 (1M context) --- AGENTS.md | 2 +- CHANGELOG.md | 9 + Project.toml | 2 - README.md | 2 +- docs/src/guide.md | 6 +- docs/src/index.md | 2 +- examples/flux_mnist/README.md | 80 ------- examples/flux_mnist/flux_mnist.jl | 93 -------- examples/flux_mnist/pytorch_mnist.py | 135 ----------- .../flux_mnist/pytorch_mnist_idiomatic.py | 151 ------------ perf/flux_cifar10/Project.toml | 14 ++ perf/flux_cifar10/README.md | 94 ++++++++ perf/flux_cifar10/flux_cifar10.jl | 201 ++++++++++++++++ perf/flux_cifar10/pytorch_cifar10.py | 224 ++++++++++++++++++ {examples => perf}/flux_mnist/Project.toml | 2 +- perf/flux_mnist/README.md | 88 +++++++ perf/flux_mnist/flux_mnist.jl | 147 ++++++++++++ perf/flux_mnist/pytorch_mnist.py | 178 ++++++++++++++ perf/{ => load_dataset}/Project.toml | 0 perf/{ => load_dataset}/README.md | 4 +- perf/{ => load_dataset}/perf.jl | 0 perf/{ => load_dataset}/perf.py | 0 src/HuggingFaceDatasets.jl | 1 - src/dataset.jl | 2 +- src/transforms.jl | 48 +++- 25 files changed, 1000 insertions(+), 485 deletions(-) delete mode 100644 examples/flux_mnist/README.md delete mode 100644 examples/flux_mnist/flux_mnist.jl delete mode 100644 examples/flux_mnist/pytorch_mnist.py delete mode 100644 examples/flux_mnist/pytorch_mnist_idiomatic.py create mode 100644 perf/flux_cifar10/Project.toml create mode 100644 perf/flux_cifar10/README.md create mode 100644 perf/flux_cifar10/flux_cifar10.jl create mode 100644 perf/flux_cifar10/pytorch_cifar10.py rename {examples => perf}/flux_mnist/Project.toml (96%) create mode 100644 perf/flux_mnist/README.md create mode 100644 perf/flux_mnist/flux_mnist.jl create mode 100644 perf/flux_mnist/pytorch_mnist.py rename perf/{ => load_dataset}/Project.toml (100%) rename perf/{ => load_dataset}/README.md (96%) rename perf/{ => load_dataset}/perf.jl (100%) rename perf/{ => load_dataset}/perf.py (100%) 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 """