Allow sampling multiple patches from the same stack

viscy/cli/cli.py

@@ -21,7 +21,6 @@ def subcommands() -> dict[str, set[str]]:
         return subcommands
 
     def add_arguments_to_parser(self, parser):
-        parser.link_arguments("data.batch_size", "model.batch_size")
         parser.link_arguments("data.yx_patch_size", "model.example_input_yx_shape")
         parser.link_arguments("model.model_config.architecture", "data.architecture")
         parser.set_defaults(

viscy/light/data.py

@@ -13,6 +13,7 @@
 from iohub.ngff import ImageArray, Plate, Position, open_ome_zarr
 from lightning.pytorch import LightningDataModule
 from monai.data import set_track_meta
+from monai.data.utils import collate_meta_tensor
 from monai.transforms import (
     CenterSpatialCropd,
     Compose,
@@ -59,9 +60,18 @@ class ChannelMap(TypedDict, total=False):
 class Sample(TypedDict, total=False):
     index: tuple[str, int, int]
     # optional
-    source: torch.Tensor
-    target: torch.Tensor
-    labels: torch.Tensor
+    source: Union[torch.Tensor, Sequence[torch.Tensor]]
+    target: Union[torch.Tensor, Sequence[torch.Tensor]]
+    labels: Union[torch.Tensor, Sequence[torch.Tensor]]
+
+
+def _collate_samples(batch: Sequence[Sample]) -> Sample:
+    elemment = batch[0]
+    collated = {}
+    for key in elemment.keys():
+        data: list[list[torch.Tensor]] = [sample[key] for sample in batch]
+        collated[key] = collate_meta_tensor([im for imgs in data for im in imgs])
+    return collated
 
 
 class NormalizeSampled(MapTransform, InvertibleTransform):
@@ -178,9 +188,15 @@ def __len__(self) -> int:
         return self._max_window
 
     def _stack_channels(
-        self, sample_images: dict[str, torch.Tensor], key: str
+        self, sample_images: list[dict[str, torch.Tensor]], key: str
     ) -> torch.Tensor:
-        return torch.stack([sample_images[ch][0] for ch in self.channels[key]])
+        if not isinstance(sample_images, list):
+            return torch.stack([sample_images[ch][0] for ch in self.channels[key]])
+        # training time
+        return [
+            torch.stack([im[ch][0] for ch in self.channels[key]])
+            for im in sample_images
+        ]
 
     def __getitem__(self, index: int) -> Sample:
         img, tz = self._find_window(index)
@@ -198,8 +214,8 @@ def __getitem__(self, index: int) -> Sample:
             sample_images["weight"] = sample_images[self.channels["target"][0]]
         if self.transform:
             sample_images = self.transform(sample_images)
-        if isinstance(sample_images, list):
-            sample_images = sample_images[0]
+        # if isinstance(sample_images, list):
+        #     sample_images = sample_images[0]
         if "weight" in sample_images:
             del sample_images["weight"]
         sample = {
@@ -275,7 +291,7 @@ class HCSDataModule(LightningDataModule):
         by default 0.8
     :param int batch_size: batch size, defaults to 16
     :param int num_workers: number of data-loading workers, defaults to 8
-    :param Literal["2.5D", "2D", "3D"] architecture: U-Net architecture,
+    :param Literal["2D", "2.1D", "2.5D", "3D"] architecture: U-Net architecture,
         defaults to "2.5D"
     :param tuple[int, int] yx_patch_size: patch size in (Y, X),
         defaults to (256, 256)
@@ -292,6 +308,8 @@ class HCSDataModule(LightningDataModule):
     :param float train_noise_std: Upper bound of the standard deviation
         of the Gaussian noise added to source images during training,
         defaults to 0.0
+    :param int train_patches_per_stack: number of patches to sample
+        from each stack during training, defaults to 1
     """
 
     def __init__(
@@ -303,14 +321,15 @@ def __init__(
         split_ratio: float = 0.8,
         batch_size: int = 16,
         num_workers: int = 8,
-        architecture: Literal["2.5D", "2D", "3D"] = "2.5D",
+        architecture: Literal["2D", "2.1D", "2.5D", "3D"] = "2.5D",
         yx_patch_size: tuple[int, int] = (256, 256),
         augment: bool = True,
         caching: bool = False,
         normalize_source: bool = False,
         ground_truth_masks: str = None,
         train_z_scale_range: tuple[float, float] = [0, 0],
         train_noise_std: float = 0.0,
+        train_patches_per_stack: int = 1,
     ):
         super().__init__()
         self.data_path = data_path
@@ -333,6 +352,12 @@ def __init__(
             raise ValueError(f"Invalid scaling range: {train_z_scale_range}")
         self.train_z_scale_range = train_z_scale_range
         self.train_noise_std = train_noise_std
+        if batch_size % train_patches_per_stack != 0:
+            raise ValueError(
+                "Batch size must be divisible by number of patches per stack. "
+                f"Got {batch_size} and {train_patches_per_stack}."
+            )
+        self.train_patches_per_stack = train_patches_per_stack
 
     def prepare_data(self):
         if not self.caching:
@@ -414,8 +439,13 @@ def _setup_fit(self, dataset_settings: dict):
         num_train_fovs = int(len(positions) * self.split_ratio)
         # training set needs to sample more Z range for augmentation
         train_dataset_settings = dataset_settings.copy()
-        expanded_z = math.ceil(self.z_window_size * (1 + self.train_z_scale_range[1]))
-        train_dataset_settings["z_window_size"] = max(1, expanded_z - expanded_z % 2)
+        z_scale_low, z_scale_high = self.train_z_scale_range
+        if z_scale_high <= 0.0:
+            expanded_z = self.z_window_size
+        else:
+            expanded_z = math.ceil(self.z_window_size * (1 + z_scale_high))
+            expanded_z -= expanded_z % 2
+        train_dataset_settings["z_window_size"] = expanded_z
         # train/val split
         self.train_dataset = SlidingWindowDataset(
             positions[:num_train_fovs],
@@ -474,10 +504,11 @@ def on_before_batch_transfer(self, batch: Sample, dataloader_idx: int) -> Sample
     def train_dataloader(self):
         return DataLoader(
             self.train_dataset,
-            batch_size=self.batch_size,
+            batch_size=self.batch_size // self.train_patches_per_stack,
             num_workers=self.num_workers,
             shuffle=True,
             persistent_workers=bool(self.num_workers),
+            collate_fn=_collate_samples,
         )
 
     def val_dataloader(self):
@@ -523,7 +554,7 @@ def _train_transform(self) -> list[Callable]:
                 keys=self.source_channel + self.target_channel,
                 w_key="weight",
                 spatial_size=(-1, self.yx_patch_size[0] * 2, self.yx_patch_size[1] * 2),
-                num_samples=1,
+                num_samples=self.train_patches_per_stack,
             )
         ]
         if self.augment:

viscy/light/engine.py

@@ -104,7 +104,6 @@ class VSUNet(LightningModule):
 
     :param dict model_config: model config,
         defaults to :py:class:`viscy.unet.utils.model.ModelDefaults25D`
-    :param int batch_size: batch size, defaults to 16
     :param Callable[[torch.Tensor, torch.Tensor], torch.Tensor] loss_function:
         loss function in training/validation, defaults to L2 (mean squared error)
     :param float lr: learning rate in training, defaults to 1e-3
@@ -127,7 +126,6 @@ class VSUNet(LightningModule):
     def __init__(
         self,
         model_config: dict = {},
-        batch_size: int = 16,
         loss_function: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] = None,
         lr: float = 1e-3,
         schedule: Literal["WarmupCosine", "Constant"] = "Constant",
@@ -145,7 +143,6 @@ def __init__(
         self.model = net_class(**model_config)
         # TODO: handle num_outputs in metrics
         # self.out_channels = self.model.terminal_block.out_filters
-        self.batch_size = batch_size
         self.loss_function = loss_function if loss_function else F.mse_loss
         self.lr = lr
         self.schedule = schedule
@@ -182,7 +179,6 @@ def training_step(self, batch: Sample, batch_idx: int):
             on_epoch=True,
             prog_bar=True,
             logger=True,
-            batch_size=self.batch_size,
             sync_dist=True,
         )
         if batch_idx < self.log_num_samples:
@@ -196,7 +192,7 @@ def validation_step(self, batch: Sample, batch_idx: int):
         target = batch["target"]
         pred = self.forward(source)
         loss = self.loss_function(pred, target)
-        self.log("loss/validate", loss, batch_size=self.batch_size, sync_dist=True)
+        self.log("loss/validate", loss, sync_dist=True)
         if batch_idx < self.log_num_samples:
             self.validation_step_outputs.append(
                 self._detach_sample((source, target, pred))

viscy/scripts/network_diagram.py

@@ -13,7 +13,6 @@
         "backbone": "convnextv2_femto",
         "stem_kernel_size": (3, 4, 4),
     },
-    batch_size=32,
 )
 # %%
 model_graph = draw_graph(