Training with PyTorch

PyTorch’s map-style dataset contract requires integer __getitem__() and __len__(). SDS datasets provide both. Use an application adapter to convert generic records into model tensors.

Install PyTorch in the training application; it is intentionally not an SDS dependency.

uv add 'torch>=2.3,<3'

import os

import signal_dataset as sds
import torch
from torch.utils.data import Dataset


class SDSTorchDataset(Dataset[dict[str, torch.Tensor]]):
    def __init__(self, uri: str) -> None:
        self._uri = uri
        self._dataset = None
        self._pid = None

    def _open(self) -> sds.Dataset:
        pid = os.getpid()
        if self._dataset is None or self._pid != pid:
            self._dataset = sds.open(self._uri)
            self._pid = pid
        return self._dataset

    def __len__(self) -> int:
        return len(self._open())

    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
        record = self._open()[index]
        return {
            "inputs": torch.from_numpy(record["samples"].data.copy()),
            "target": torch.tensor(record.metadata["label"], dtype=torch.long),
        }

    def __getstate__(self) -> dict[str, object]:
        return {"_uri": self._uri}

    def __setstate__(self, state: dict[str, object]) -> None:
        self._uri = str(state["_uri"])
        self._dataset = None
        self._pid = None

Field.data is immutable. The example copies it before torch.from_numpy so model transforms cannot mutate the SDS record or trigger PyTorch’s non-writable-array warning.

Use PyTorch-owned sampling and batching:

from torch.utils.data import DataLoader

dataset = SDSTorchDataset("captures.sds")
loader = DataLoader(
    dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=worker_count,
    pin_memory=use_accelerator,
)

for batch in loader:
    train_step(batch["inputs"], batch["target"])

The PID check reopens the same immutable SDS URI in both forked and spawned workers. Pickling also drops the live reader. Return CPU tensors from workers; let pin_memory and the training step manage accelerator transfer.

For distributed training, let DistributedSampler assign ordinals. Call set_epoch() before each epoch so its shuffle order changes:

from torch.utils.data.distributed import DistributedSampler

sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler)

for epoch in range(num_epochs):
    sampler.set_epoch(epoch)
    for batch in loader:
        train_step(batch["inputs"], batch["target"])

Keep model conversion, augmentation, sampling, and batching in the training application. SDS remains responsible only for immutable indexed records and aligned annotations.

Batched fetch

DataLoader calls __getitems__ on a map-style dataset when one exists (since PyTorch 1.13), passing the whole batch of indices at once. That is worth implementing here, because the underlying container is markedly faster reading a list of indices in one call than the same indices one at a time:

    def __getitems__(self, indices: list[int]) -> list[dict[str, torch.Tensor]]:
        return [self[index] for index in indices]

That version still reads one at a time. To get the batched read as well, address the shards yourself with dataset.read_instructions(...) and fetch each run in one call. See bring your own loader. Note that Grain does not call __getitems__; it looks for _getitems behind a separate opt-in protocol, so a source used by both needs the method under both names.