From ac1144ed5dda661b04cf172b91669eef2a124a34 Mon Sep 17 00:00:00 2001 From: William Patton Date: Wed, 13 Nov 2024 10:49:06 -0800 Subject: [PATCH 1/6] fix bug in cnnectome_unet return types The Architecture superclass says that we should be returning coordinates for these properties This is important because we regularly use these values in arithmetic where we expect to execute +-*/ element wise --- dacapo/experiments/architectures/cnnectome_unet.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/dacapo/experiments/architectures/cnnectome_unet.py b/dacapo/experiments/architectures/cnnectome_unet.py index d89e902ac..f706c4fdb 100644 --- a/dacapo/experiments/architectures/cnnectome_unet.py +++ b/dacapo/experiments/architectures/cnnectome_unet.py @@ -3,6 +3,8 @@ import torch import torch.nn as nn +from funlib.geometry import Coordinate + import math @@ -176,7 +178,7 @@ def __init__(self, architecture_config): self.unet = self.module() @property - def eval_shape_increase(self): + def eval_shape_increase(self) -> Coordinate: """ The increase in shape due to the U-Net. @@ -192,7 +194,7 @@ def eval_shape_increase(self): """ if self._eval_shape_increase is None: return super().eval_shape_increase - return self._eval_shape_increase + return Coordinate(self._eval_shape_increase) def module(self): """ @@ -306,11 +308,11 @@ def scale(self, voxel_size): The voxel size should be given as a tuple ``(z, y, x)``. """ for upsample_factor in self.upsample_factors: - voxel_size = voxel_size / upsample_factor + voxel_size = voxel_size / Coordinate(upsample_factor) return voxel_size @property - def input_shape(self): + def input_shape(self) -> Coordinate: """ Return the input shape of the U-Net. From 08d7074ebe6c5b4d4a7073bec455a2b6d40cc6a7 Mon Sep 17 00:00:00 2001 From: William Patton Date: Wed, 13 Nov 2024 10:49:53 -0800 Subject: [PATCH 2/6] CNNectomeUNet fix kernel size logic --- dacapo/experiments/architectures/cnnectome_unet.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/dacapo/experiments/architectures/cnnectome_unet.py b/dacapo/experiments/architectures/cnnectome_unet.py index f706c4fdb..116be284c 100644 --- a/dacapo/experiments/architectures/cnnectome_unet.py +++ b/dacapo/experiments/architectures/cnnectome_unet.py @@ -237,16 +237,15 @@ def module(self): """ fmaps_in = self.fmaps_in levels = len(self.downsample_factors) + 1 - dims = len(self.downsample_factors[0]) - if hasattr(self, "kernel_size_down"): + if self.kernel_size_down is not None: kernel_size_down = self.kernel_size_down else: - kernel_size_down = [[(3,) * dims, (3,) * dims]] * levels - if hasattr(self, "kernel_size_up"): + kernel_size_down = [[(3,) * self.dims, (3,) * self.dims]] * levels + if self.kernel_size_up is not None: kernel_size_up = self.kernel_size_up else: - kernel_size_up = [[(3,) * dims, (3,) * dims]] * (levels - 1) + kernel_size_up = [[(3,) * self.dims, (3,) * self.dims]] * (levels - 1) # downsample factors has to be a list of tuples downsample_factors = [tuple(x) for x in self.downsample_factors] @@ -326,7 +325,7 @@ def input_shape(self) -> Coordinate: Note: The input shape should be given as a tuple ``(batch, channels, [length,] depth, height, width)``. """ - return self._input_shape + return Coordinate(self._input_shape) @property def num_in_channels(self) -> int: From a0b51165d56647d77fb05d3fd7b471aec279e332 Mon Sep 17 00:00:00 2001 From: William Patton Date: Wed, 13 Nov 2024 10:52:52 -0800 Subject: [PATCH 3/6] CNNectomeUNet: make the final conv pass in the upsample pass a bit more robust If we upsample, we probably want to apply a convolution to finetune the outputs rather than simply upsampling which we could do outside of a network. If we assume a kernel of size (3, 3, 3), it fails for 2D networks that process using kernels of size (1, 3, 3). We now just use the last kernel in the kernel size up. This is a bit more robust. --- dacapo/experiments/architectures/cnnectome_unet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dacapo/experiments/architectures/cnnectome_unet.py b/dacapo/experiments/architectures/cnnectome_unet.py index 116be284c..c064305b1 100644 --- a/dacapo/experiments/architectures/cnnectome_unet.py +++ b/dacapo/experiments/architectures/cnnectome_unet.py @@ -281,7 +281,7 @@ def module(self): conv = ConvPass( self.fmaps_out, self.fmaps_out, - [(3,) * len(upsample_factor)] * 2, + kernel_size_up[-1], activation="ReLU", batch_norm=self.batch_norm, ) From 36cbd6724b793810c7306cc5cec35f064617e78c Mon Sep 17 00:00:00 2001 From: William Patton Date: Wed, 13 Nov 2024 10:53:37 -0800 Subject: [PATCH 4/6] Gunpowder Trainer: if the raw data doesn't have a channel dim, add it during training Otherwise the BatchNorm breaks --- dacapo/experiments/trainers/gunpowder_trainer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dacapo/experiments/trainers/gunpowder_trainer.py b/dacapo/experiments/trainers/gunpowder_trainer.py index dcb40c115..bea9c96e2 100644 --- a/dacapo/experiments/trainers/gunpowder_trainer.py +++ b/dacapo/experiments/trainers/gunpowder_trainer.py @@ -173,6 +173,8 @@ def build_batch_provider(self, datasets, model, task, snapshot_container=None): assert isinstance(dataset.weight, int), dataset raw_source = gp.ArraySource(raw_key, dataset.raw) + if dataset.raw.channel_dims == 0: + raw_source += gp.Unsqueeze([raw_key], axis=0) if self.clip_raw: raw_source += gp.Crop( raw_key, dataset.gt.roi.snap_to_grid(dataset.raw.voxel_size) From 23d1c22cf83a47d6ff735c43a48a097b89d94330 Mon Sep 17 00:00:00 2001 From: William Patton Date: Wed, 13 Nov 2024 10:54:20 -0800 Subject: [PATCH 5/6] Datasplit test fixture: use the voxel_size attribute for voxel size. This should probably just be switched to use `funlib.perisistence.prepare_ds` --- tests/fixtures/datasplits.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/fixtures/datasplits.py b/tests/fixtures/datasplits.py index 448c9c834..73c282a89 100644 --- a/tests/fixtures/datasplits.py +++ b/tests/fixtures/datasplits.py @@ -73,10 +73,10 @@ def twelve_class_datasplit(tmp_path): gt_dataset[:] += random_data > i raw_dataset[:] = random_data raw_dataset.attrs["offset"] = (0, 0, 0) - raw_dataset.attrs["resolution"] = (2, 2, 2) + raw_dataset.attrs["voxel_size"] = (2, 2, 2) raw_dataset.attrs["axis_names"] = ("z", "y", "x") gt_dataset.attrs["offset"] = (0, 0, 0) - gt_dataset.attrs["resolution"] = (2, 2, 2) + gt_dataset.attrs["voxel_size"] = (2, 2, 2) gt_dataset.attrs["axis_names"] = ("z", "y", "x") crop1 = RawGTDatasetConfig(name="crop1", raw_config=crop1_raw, gt_config=crop1_gt) @@ -184,10 +184,10 @@ def six_class_datasplit(tmp_path): gt_dataset[:] += random_data > i raw_dataset[:] = random_data raw_dataset.attrs["offset"] = (0, 0, 0) - raw_dataset.attrs["resolution"] = (2, 2, 2) + raw_dataset.attrs["voxel_size"] = (2, 2, 2) raw_dataset.attrs["axis_names"] = ("z", "y", "x") gt_dataset.attrs["offset"] = (0, 0, 0) - gt_dataset.attrs["resolution"] = (2, 2, 2) + gt_dataset.attrs["voxel_size"] = (2, 2, 2) gt_dataset.attrs["axis_names"] = ("z", "y", "x") crop1 = RawGTDatasetConfig( From 270c4692864142dd084100930051dd096bdd2230 Mon Sep 17 00:00:00 2001 From: William Patton Date: Wed, 13 Nov 2024 10:55:12 -0800 Subject: [PATCH 6/6] 2D model still only gets data with a single channel from the trainer --- tests/operations/test_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/operations/test_train.py b/tests/operations/test_train.py index be0a94d16..e32276cec 100644 --- a/tests/operations/test_train.py +++ b/tests/operations/test_train.py @@ -49,7 +49,7 @@ def unet_architecture(batch_norm, upsample, use_attention, three_d): name=name, input_shape=(2, 132, 132), eval_shape_increase=(8, 32, 32), - fmaps_in=2, + fmaps_in=1, num_fmaps=8, fmaps_out=8, fmap_inc_factor=2,