Record assembly and batched loading

WarningContent review status

This page contains unreviewed content created with large language model (LLM) tools.

A tensorbids record set joins image references and relation columns on BIDS entity keys. Record functions can express array indexing and spatial windows lazily. The bitsio resolver converts those selections to tensor-store region reads, and the PyTorch-compatible loader schedules and collates them in batches.

1 Definitions

A BIDSReference selects one image for each record and defines the requested representation:

  • the return type is TENSOR, METATENSOR, OBJECT, or BYTES;
  • the port selects the reconstructed image tensor or a named navigator port, including "stored"; and
  • dequantize="attach" returns integer codes with their scale and offset as a Codes value instead of applying the scale during the read.

A relation field is addressed as <relation>.<column>. A bare column name is accepted when it is unambiguous among the selected relations.

The record key determines the unit of observation. By default, tensorbids derives it from the BIDS entities that vary among the selected references and includes subject when present. The optional key argument overrides that choice. Each assembled record must contain one value for every reference and field; missing="drop" removes incomplete records, while missing="error" rejects them.

from bitsio import BITSZarrDataset
from tensorbids.query import BIDSReference

ds = BITSZarrDataset.open("ds000114.bits")
print(ds.records(BIDSReference(suffix="T1w")).key)
print(ds.records(BIDSReference(suffix="bold")).key)
print(len(ds.records(BIDSReference(suffix="T1w"), BIDSReference(suffix="bold"))))
('subject',)
('subject', 'session')
2

The T1-weighted selection varies by subject, whereas the BOLD selection varies by subject and session. Their combined key is therefore (subject, session). Only subject/session pairs that contain both image types remain.

2 Lazy reads

Inside a record function, tensor references are represented by LazyTensor objects. Basic indexing, random_window(), and bound_to_mask() update a read plan without loading data. When a batch is evaluated, compatible operations are reduced to rectangular array regions and submitted together. The tensor store reads the chunks intersecting those regions; see read amplification.

Random-window sampling derives its generator state from the record-set seed, epoch, record index, and reference name. A fixed seed and epoch therefore produce the same windows across worker processes. seed=None requests non-deterministic draws.

def crop(record):
    return record["T1w"].random_window(i=32, j=32, k=32)


windows = ds.records(BIDSReference(suffix="T1w")).select(crop, seed=0)
first, again = windows.read([0]), windows.read([0])
windows.set_epoch(1)
print((first[0] == again[0]).all(), (first[0] == windows.read([0])[0]).all())
True False

With stack=True, values for each reference are stacked along a batch axis. All arrays must have a common shape; fixed-size windows satisfy this condition even when their source images differ. OBJECT and BYTES results are not stacked.

3 The loader

loader() returns a PyTorch IterableDataset whose items are already complete batches. Wrap it with torch.utils.data.DataLoader(batch_size=None) to avoid a second batching layer. Workers partition the epoch’s batch order. in_flight limits outstanding reads, while slots controls reusable shared-memory buffers. set_epoch() updates both shuffle order and deterministic random windows.

With the recommended forkserver or spawn process start method, a worker starts a fresh interpreter and imports the record function by module and name. Define the function in an importable module and protect executable script code with if __name__ == "__main__":. Record sets backed by workspace views are not currently serializable to worker processes and must use num_workers=0.

4 Compared with MONAI

The record-set interface is defined in tensorbids and can exploit bitsio region reads before data reaches the application. MONAI (Cardoso et al. 2022) provides a broader transform ecosystem around item dictionaries and MetaTensor. bitsfuse.monai.BitsDataset adapts BITS records to that interface, generally reading the selected image before applying MONAI transforms. See Use BITS data in a MONAI pipeline.

References

Cardoso, M. Jorge, Wenqi Li, Richard Brown, et al. 2022. MONAI: An Open-Source Framework for Deep Learning in Healthcare. https://arxiv.org/abs/2211.02701.