Build records and batches

WarningContent review status

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

Join one image per subject with fields from tabular relations, apply a lazy spatial crop, and iterate over batches suitable for a model.

1 Add a table

Create a relation that assigns each subject to one of two folds. Reopen the dataset in writable mode before adding the relation:

import polars as pl

from bitsio import BITSZarrDataset

ds = BITSZarrDataset.open("ds000114.bits", read_only=False)
subjects = ds.subjects()
folds = pl.DataFrame({
    "subject": subjects,
    "fold": [n % 2 for n in range(len(subjects))],
})
written = ds.write_table("split/folds", folds)  # idempotent
print(written.name, written.write.dtypes)
split/folds {'subject': 'string', 'fold': 'int64'}

The registered name "split/folds" identifies this relation as a subject split named folds. write_table() accepts an in-memory data frame or a path to a CSV or TSV file. For an in-memory frame, the writer preserves the declared column types.

2 Build a record set

A BITS record set joins image references and relation columns at a shared key. The following record set contains one T1-weighted image, the participant’s dominant hand, and the assigned fold for each subject:

from tensorbids.query import BIDSReference

records = ds.records(
    BIDSReference(suffix="T1w"),
    "participants.dominant_hand",
    "split/folds.fold",
)
print(records)
columns = ("subject", "participants.dominant_hand", "split/folds.fold")
print(records.table.select(columns).head(4))
print(len(records.where(fold=0)))
BITSRecordSet(10 records, key=['subject'], references=['T1w'], fields=['participants.dominant_hand', 'split/folds.fold'])
shape: (4, 3)
┌─────────┬────────────────────────────┬──────────────────┐
│ subject ┆ participants.dominant_hand ┆ split/folds.fold │
│ ---     ┆ ---                        ┆ ---              │
│ str     ┆ str                        ┆ i64              │
╞═════════╪════════════════════════════╪══════════════════╡
│ 01      ┆ left                       ┆ 0                │
│ 02      ┆ right                      ┆ 1                │
│ 03      ┆ right                      ┆ 0                │
│ 04      ┆ right                      ┆ 1                │
└─────────┴────────────────────────────┴──────────────────┘
5

BIDSReference selects image rows from the BIDS index. A string of the form "relation.column" selects a relation column. The join key is inferred from the selected data; this example produces one record per subject. Records with a missing component are dropped by default. Pass missing="error" to reject incomplete joins. where(fold=0) filters the assembled records.

3 Batching data for models

The T1-weighted images do not all have the same shape, so they cannot be stacked directly. A record function requests a 64 x 64 x 64 window from each image and renames the output fields. The window remains lazy until the loader resolves the batch, allowing the backend to read only the required array region.

crops.py
def crop(record):
    """A random 64-voxel cube of the T1w image, and the label."""
    return {
        "T1w": record["T1w"].random_window(i=64, j=64, k=64),
        "hand": record["participants.dominant_hand"],
    }
WarningData loader workers

With the recommended forkserver or spawn process start method, each worker imports the record function by module and name. Define the function in an importable module, rather than as a notebook-local function or lambda.

import torch

from crops import crop

loader = records.select(crop, seed=0).loader(batch_size=2, shuffle=True, stack=True)
data = torch.utils.data.DataLoader(
    loader,
    batch_size=None,
    num_workers=2,
    persistent_workers=True,
    multiprocessing_context="forkserver",
)
for epoch in range(2):
    loader.set_epoch(epoch)
    batch = next(iter(data))
    t1w = batch["T1w"]
    print(epoch, type(t1w).__name__, tuple(t1w.shape), t1w.dtype, batch["hand"])
0 Tensor (2, 64, 64, 64) torch.float32 ['left', 'left']
1 Tensor (2, 64, 64, 64) torch.float32 ['right', 'left']

records.loader() already produces complete batches, so the outer PyTorch DataLoader uses batch_size=None. set_epoch() changes the deterministic shuffle order and random-window draws for the next epoch in every worker.

WarningWorkers and the main module

In a script, protect the training entry point with if __name__ == "__main__": so that worker imports do not run the training loop.

NoteBITS and MONAI

The record-set loader is a tensorbids data-loading interface. The temporary MONAI integration is maintained by bitsfuse; see Use BITS data in a MONAI pipeline.

4 Where to next