rfgen.adapters

Optional framework adapters that bridge rfgen’s consumer abstractions (BaseSampler, BaseCollator) to framework-specific training APIs.

The core rfgen.storage and rfgen.dataset modules do not depend on any training framework. Only the adapter modules introduce framework dependencies, and each is behind a rfgen[<extra>] install extra.

Warning

Pre-implementation. This page describes proposed contracts. Class signatures, parameter types, config field names, and behavior are subject to change before code lands. Once implementation exists, content here will be regenerated from docstrings or sourced from running tests.

Class index

Class	Kind	Module	Notes
RfgenTorchDataset	concrete	rfgen.dataset	IterableDataset wrapper for torch.utils.data.DataLoader; behind rfgen[torch]

---

class rfgen.dataset.RfgenTorchDataset

from torch.utils.data import IterableDataset

class RfgenTorchDataset(IterableDataset):
    """Wrap a rfgen BaseStore so it can be passed to torch.utils.data.DataLoader.

    Behind the rfgen[torch] extra. Importing rfgen.dataset does not import
    torch.utils.data; this optional class is loaded lazily.
    """

    def __init__(
        self,
        store: BaseStore,
        *,
        sampler: BaseSampler | None = None,
        collator: BaseCollator | None = None,
    ): ...

    def __iter__(self) -> Iterator[Record | dict[str, torch.Tensor]]:
        """Yields collated batches if collator is set, else raw Records."""

Warning

Pre-implementation. This page describes proposed contracts. Class signatures, parameter types, config field names, and behavior are subject to change before code lands. Once implementation exists, content here will be regenerated from docstrings or sourced from running tests.

Wraps a BaseStore as a torch.utils.data.IterableDataset so it can be passed directly to torch.utils.data.DataLoader.

Each DataLoader worker opens its own StoreHandle from the store’s dataset_uri(). The adapter is loaded lazily from rfgen.dataset, so importing the rest of the consumer toolkit remains framework-neutral.

When collator is set, __iter__ yields dict[str, torch.Tensor] batches shaped by the collator. When collator=None, __iter__ yields raw Record objects.

When sampler is set, iteration order follows the sampler. When sampler=None, a SequentialSampler is used by default.

Constructor parameters

Name	Type	Default	Description
store	BaseStore	-	Store whose dataset_uri() each worker opens for reading
sampler	BaseSampler | None	None	Iteration order; defaults to SequentialSampler
collator	BaseCollator | None	None	Batch shaping; when None, yields raw Record objects

Notes

• Framework-neutral core. rfgen.storage and the default rfgen.dataset import path do not import torch.utils.data. The optional RfgenTorchDataset symbol is loaded lazily behind rfgen[torch].

• DataLoader compatibility. Because RfgenTorchDataset is an IterableDataset, DataLoader sharding for multi-worker loading is controlled by torch.utils.data.get_worker_info(). When num_workers > 1, the __iter__ implementation requires a sampler with a concrete partition(worker_id, num_workers) hook using worker_info.id and worker_info.num_workers. Duck-typed samplers that provide only __iter__ and __len__ are supported only with num_workers <= 1; multi-worker custom samplers must implement partition(...) or bind(source) that returns a sampler with a concrete partition(...) implementation.

• Collator and batch_size. When collator is set, DataLoader should be constructed with batch_size=None (or batch_size=1 with collate_fn disabled) because RfgenTorchDataset.__iter__ already returns batched tensors from the collator. When collator=None, let DataLoader’s default collation handle stacking.

Usage example

from rfgen.storage import ZarrStore
from rfgen.dataset.samplers import RandomSampler
from rfgen.dataset.collators import DetectionCollator
from rfgen.dataset import RfgenTorchDataset
from torch.utils.data import DataLoader

store = ZarrStore()
sampler = RandomSampler(seed=42)
collator = DetectionCollator()
dataset = RfgenTorchDataset(store, sampler=sampler, collator=collator)
loader = DataLoader(dataset, batch_size=None, num_workers=4)
for batch in loader:
    iq = batch["iq"]          # shape (1, 2, N)
    bboxes = batch["bboxes"]  # list[Tensor], one entry for this record
    ...

---

See Also

• Concepts / Architecture § Consumer Access - framing of the adapter layer.

• Reference / rfgen.dataset - sampler and collator ABCs and concrete classes.

• Reference / rfgen.storage - StoreHandle API.