cellmap_flow.finetune.data.loader

The DataLoaders the finetune trainer reads its patches through.

create_dataloader builds one from a session’s corrections directory (its _virtual_sources.json manifest); make_training_loader and rebuild_loader make sure every loader the trainer uses, the OOM fallback’s included, keeps its workers from one epoch to the next.

Attributes

`logger`
`PatchVoxels`

Functions

`training_patch_voxels`(→ Optional[PatchVoxels])	The (input, output) patch in voxels `model_config` trains on, or
`dataset_from_manifest`(...)	Instantiate a `VirtualPatchDataset` from a manifest dict.
`create_dataloader`(→ torch.utils.data.DataLoader)	Build the training DataLoader for a corrections directory.
`make_training_loader`(→ torch.utils.data.DataLoader)	The DataLoader the trainer reads patches through.
`rebuild_loader`(→ torch.utils.data.DataLoader)	`loader` again with a different batch size and everything else kept.

Module Contents

cellmap_flow.finetune.data.loader.logger

cellmap_flow.finetune.data.loader.PatchVoxels

cellmap_flow.finetune.data.loader.training_patch_voxels(model_config) → PatchVoxels | None

The (input, output) patch in voxels model_config trains on, or None for the session’s (the model’s read and write shapes).

A model type whose network takes a fixed tile says so with a training_patch_voxels() method (Cellpose-SAM: one 256 x 256 slice). A finetuned model trains as its base does, so a base_model_config is asked in turn.

cellmap_flow.finetune.data.loader.dataset_from_manifest(manifest: dict, corrections_dir: str | None = None, augment: bool | None = None, patch_voxels: PatchVoxels | None = None) → cellmap_flow.finetune.data.dataset.VirtualPatchDataset

Instantiate a VirtualPatchDataset from a manifest dict.

Recognized manifest kinds:

volume_zarr_v1 (current): trainer reads the session’s annotation_volume.zarr directly. Field volume_zarr_path.

corrections_dir is where the session’s good regions are looked up; omit it and the dataset simply trains without rehearsal anchors.

patch_voxels, (input, output) in voxels, replaces the manifest’s patch sizes, which are the model’s serving geometry: for a model whose network trains on another (training_patch_voxels). The voxel sizes stay the manifest’s. An epoch then counts as many more patches per annotated chunk as the patch is smaller, so that it still covers what was annotated.

cellmap_flow.finetune.data.loader.create_dataloader(corrections_zarr_path: str, batch_size: int = 2, augment: bool = True, num_workers: int = 4, patch_voxels: PatchVoxels | None = None) → torch.utils.data.DataLoader

Build the training DataLoader for a corrections directory.

Requires a _virtual_sources.json manifest. Every path that creates a session writes one (volume creation, YAML import, training submit) and restarts backfill one, so a missing manifest means something upstream failed – which is worth an exception rather than training on anything else. The patch geometry comes from the manifest, and the dataset samples randomly, so there is no patch shape or shuffle to pass.

Parameters:

corrections_zarr_path – Session corrections directory.
batch_size – Clamped down to the dataset size when smaller.
augment – Flips, XY rotations, brightness and noise, on top of the patch-centre jitter that is always applied. It overrides the manifest’s stored preference.
num_workers – DataLoader workers. Spawned, not forked – tensorstore handles do not survive fork.
patch_voxels – (input, output) patch in voxels instead of the manifest’s, for a model that trains on its own (training_patch_voxels).

cellmap_flow.finetune.data.loader.make_training_loader(dataset, batch_size: int, num_workers: int) → torch.utils.data.DataLoader

The DataLoader the trainer reads patches through.

Persistent, spawned workers: each keeps its copy of the dataset, and so its RNG, from one epoch to the next. Non-persistent workers are re-spawned every epoch from a fresh pickle of the dataset – RNG unset, same seed – so every epoch would draw the identical patches and augmentations.

cellmap_flow.finetune.data.loader.rebuild_loader(loader: torch.utils.data.DataLoader, batch_size: int) → torch.utils.data.DataLoader

loader again with a different batch size and everything else kept.

The OOM fallback in the trainer used to rebuild its loader with only some of the original arguments; losing persistent_workers made every epoch after an OOM repeat the same patches (see make_training_loader).