netcl.data — DataLoader, Augment, SHM-Ring

The data API is the training-data ingestion layer. It builds on top of Python's multiprocessing to give you a DataLoader with a worker pool, a sliding-window prefetch ring buffer, and an optional lock-free shared-memory fast path for the inner H2D copy on OpenCL devices. It also ships a small set of CPU and GPU augmentation primitives, balanced-sampling / label filters, and a fork-vs-spawn aware pool initializer that keeps PyOpenCL safe to import inside a worker.

Note — Long-form imports. The netcl/data package does not have a flat __init__.py re-export. Reach the public symbols by their full submodule path:

python from netcl.data.dataloader import DataLoader from netcl.data.gpu_batch import GpuBatch, to_gpu_batch from netcl.data.shared_batch import SharedBatchRingBuffer, worker_write_to_shm from netcl.data.augment import random_erase, cutout, color_jitter from netcl.data.augment_gpu import flip_horizontal, apply_color_jitter, cutout from netcl.data.filters import ( apply_filters, to_float, normalize, horizontal_flip, random_crop, default_cifar10_filters, basic_cifar10_filters, )

Module Map

DataLoader

DataLoader is the only thing most user code should instantiate. It combines a worker pool, an asynchronous prefetch ring, and a thin post-processing hook chain. It is lazy: the worker pool is only created on the first call to __iter__, and it is persistent: the same pool is reused across epochs as long as the dataset identity is unchanged.

from netcl.data.dataloader import DataLoader

loader = DataLoader(
    dataset,
    batch_size=128,
    shuffle=True,
    drop_last=False,
    num_workers=2,
    prefetch=4,            # ring buffer depth (defaults to 2 * num_workers)
    seed=0,
    transforms=None,       # optional (xb, yb) -> (xb, yb) callable
    use_shared_mem=False,  # set True for the SHM fast path (num_workers must be > 0)
)
for xb, yb in loader:      # numpy arrays by default — see bind_queue() below
    ...

DataLoader exposes three lifecycle hooks beyond __iter__:

• close() — terminate the worker pool and unlink() the SHM ring buffer. Called automatically from __del__, but explicit calls are encouraged in long-running processes to free the SHM segments deterministically.
• __len__() — the number of batches the loader will yield this epoch (respects drop_last).
• __del__() — calls close(). Safe to rely on for short-lived scripts.

Fork-vs-Spawn Caveat

PyOpenCL is fork-unsafe: a fork() after PyOpenCL has imported its ICD / bound a driver leaves the child process with a half-initialized runtime that crashes on the first kernel launch. The DataLoader handles this automatically:

• On Linux, where fork is the default and cheap (it lets workers inherit the parent's loaded datasets via copy-on-write), _get_mp_context() returns multiprocessing.get_context("fork"). PyOpenCL is not yet imported in the main process when the workers start, so the fork happens before any driver binding.
• On macOS and Windows, the only available start method is "spawn". The dataset is sent to each worker once via the pool initializer; the first iteration may take a second or two for a large in-memory dataset, but every subsequent iteration is fast.

If you need to use fork explicitly (e.g. for a one-shot diagnostic) and PyOpenCL is already imported, do it before the DataLoader ever constructs its pool. Otherwise stick with the platform default.

Shared-Memory Ring Buffer

When use_shared_mem=True and num_workers > 0, DataLoader takes the zero-copy path. The main process creates a SharedBatchRingBuffer sized K = prefetch + 2 and asks each worker to write its batch directly into a slot instead of returning a pickled NumPy tuple.

loader = DataLoader(ds, batch_size=64, num_workers=2, prefetch=8, use_shared_mem=True)

The flow is:

1. The first batch is processed on the main thread (so the loader can learn the field shapes and dtypes), and yielded to the consumer immediately.
2. The main process attaches to the SHM ring and starts apply_async jobs on the worker pool, each call passing the slot ID as a serialized attach_info dict (name, slot size, shape, dtype).
3. Workers call worker_write_to_shm which np.copytos the stacked arrays into the slot.
4. The main process slot_copy(slot_id)s the data out, freeing the slot for the next worker. The slot's view is then np.array(...) copied into a fresh allocation, so the SHM slot can be reused immediately without a race.
5. The last (incomplete) batch falls back to the regular pickle path because the per-worker shape discovery happens only once.

The full shared-memory contract — including the per-worker cache that avoids re-attaching on every call — is described in Distributed Architecture.

GPU Prefetch: bind_queue() and GpuBatch

DataLoader does not expose a pin_memory=True constructor flag; instead it exposes GPU prefetch directly, which does strictly more than pinning the host buffer — it uploads the batch to the device during the load phase, overlapping the H2D copy with the previous step's GPU work instead of doing it on the training loop's critical path.

from netcl.data.dataloader import DataLoader

loader = DataLoader(dataset, batch_size=128, num_workers=4, use_shared_mem=True)
loader.bind_queue(queue)   # call once, after GPU init

for batch in loader:
    xb, yb = batch[0], batch[1]   # xb is already a device Tensor
    ...

Call bind_queue(queue) once, after the OpenCL device/queue exists. From then on every batch yielded by __iter__ is a netcl.data.gpu_batch.GpuBatch — a tuple subclass whose float32, C-contiguous fields have already been converted to pooled device Tensors via a non-blocking cl.enqueue_copy. Fields of any other dtype (integer labels, for example) are left as plain NumPy arrays. On the OpenCL in-order queue, any kernel you enqueue afterwards (the forward pass) is automatically ordered behind the pending copies — you never need to call wait() yourself.

from netcl.data.gpu_batch import GpuBatch

batch = next(iter(loader))
assert isinstance(batch, GpuBatch)
batch.h2d_done()             # True once the copies have completed on-device
batch.h2d_done(wait=True)    # block until they have (rarely needed)

Why this matters: converting a batch with Tensor.from_host(...) per-field on the training loop's main thread routes through a synchronous driver-side memcpy for a fresh buffer allocation — for a ~44 MB batch this was measured at ~50 ms of pure main-thread time. bind_queue() moves that entirely off the critical path: the pooled-buffer, non-blocking copy costs well under 1 ms of host time per batch, because the actual DMA runs concurrently with the GPU while the next batch is being loaded.

If the GPU upload raises (OOM, a lost context, …), _postprocess catches the exception and falls back to the plain NumPy batch for that iteration rather than propagating the error — training does not crash because of a transient H2D failure.

When use_shared_mem=True, the SHM ring slot backing a GpuBatch is not returned to the free pool immediately after the batch is yielded — a worker must not overwrite the slot while the async copy is still reading it. The loader defers recycling until batch.h2d_done() reports completion (checked opportunistically after each subsequent iteration, and awaited at epoch end), which is normally instantaneous by the time the next few batches have been produced.

Legacy Constructor Arguments

device_queue, augment, overlap, autocast, and async_transfer predate bind_queue() and remain for backward compatibility:

New code should prefer bind_queue() over device_queue=..., async_transfer=True — they do the same thing, but bind_queue() can be called after construction (important on fork platforms, where the loader/worker pool must be created before the OpenCL context exists, but the queue itself only exists after).

pin_memory-Equivalent Single-Tensor Uploads

Outside the DataLoader, the same pinned-vs-async choice is available per-tensor via Tensor.from_host(...): use_pinned=True (the default when NETCL_PINNED_H2D=1) routes the H2D copy through a page-locked staging buffer in the PinnedBufferPool, and async_copy=True (the default when NETCL_ASYNC_H2D=1) makes the copy non-blocking. When both are requested, the async path takes priority (a blocking pinned-buffer map on a busy in-order queue would stall the host on the entire queued pipeline) — see Tensor for the full precedence rule.

Augmentations

CPU Augmentations

from netcl.data.augment import random_erase, cutout, color_jitter

# All functions take NCHW float32 (or uint8) numpy arrays and return a fresh array.
x_aug = random_erase(x, p=0.5, scale=(0.02, 0.2), ratio=(0.3, 3.3))
x_aug = cutout(x, size=8)
x_aug = color_jitter(x, brightness=0.2, contrast=0.2, saturation=0.0)

The CPU augmentations are intentionally tiny — they exist to make single-process debug runs easy. In a real training loop you should reach for the GPU variants, or for the filter pipeline below.

GPU Augmentations

from netcl.data.augment_gpu import flip_horizontal, apply_color_jitter, cutout

# All functions take an OpenCL Tensor and per-sample masks/factors from the host.
mask = (np.random.rand(N) < 0.5).astype(np.uint8)
x_aug = flip_horizontal(x, mask)

b = np.random.uniform(-0.2, 0.2, size=(N,)).astype(np.float32)
c = 1.0 + np.random.uniform(-0.2, 0.2, size=(N,)).astype(np.float32)
x_aug = apply_color_jitter(x, b, c)

centers = np.random.randint(0, 32, size=(N, 2)).astype(np.int32)
x_aug = cutout(x, centers, size=8)

Note — Randomness is host-side. Every GPU augment function takes its randomness (masks, factors, centers) as a host NumPy array. The OpenCL kernel stays a pure deterministic function of (x, randomness), which is what makes these kernels trivially testable and re-buildable across devices.

Filter Pipeline

The data/filters.py module provides a composable filter pipeline that runs on the dataset (not the batch) and is therefore most useful when combined with the DataLoader.transforms= argument.

from netcl.data.filters import (
    random_crop, horizontal_flip, to_float, normalize,
    default_cifar10_filters, basic_cifar10_filters,
)

pipeline = [
    random_crop(padding=4, crop_size=32, pad_mode="reflect"),
    horizontal_flip(p=0.5),
    to_float(scale=255.0),
    normalize(mean=(0.4914, 0.4822, 0.4465), std=(0.2023, 0.1994, 0.2010)),
]

loader = DataLoader(ds, batch_size=128, transforms=pipeline)

A FilterFn is just Callable[[np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray]]. You can write your own — they compose left-to-right in the order they appear in the list.

Custom Datasets

Any object that supports __len__ and __getitem__ works. A sample is a tuple of NumPy arrays; a batch is what you get when the loader calls np.stack along a new leading axis:

class MyDataset:
    def __len__(self):
        return 1000

    def __getitem__(self, i):
        x = ...           # np.ndarray, e.g. (C, H, W)
        y = ...           # np.ndarray, e.g. ()  (scalar)
        return x, y       # tuple of ndarrays

For multi-field samples, return a longer tuple (or a dict with the same keys in every item). The loader will np.stack each field along a new leading axis and yield them in the same shape:

return x_img, x_mask, y_label
# loader yields (xb_img, xb_mask, yb_label) with shapes ((B, C, H, W), (B, H, W), (B,))

See also

• MNIST with MLP — the worked example that uses DataLoader end-to-end.
• Tensor — the device-side value type GpuBatch fields become, and the from_host(async_copy=, use_pinned=) precedence rule bind_queue() builds on.
• OpenCLBackend — the PinnedBufferPool that backs the single-tensor pinned H2D path.
• Memory Pool — the persistent pool to_gpu_batch() allocates device buffers from, and how to size it.
• Distributed Architecture — the cross-process IPC contract that the SharedBatchRingBuffer implements.