From 2d4d83ec80bfb55d3a60fb7863db2185ea393905 Mon Sep 17 00:00:00 2001 From: Benjamin Morris Date: Tue, 17 Jun 2025 14:55:06 -0700 Subject: [PATCH] renaming --- cyto_dl/nn/vits/decoder.py | 5 ++++- cyto_dl/nn/vits/mae.py | 1 + cyto_dl/nn/vits/seg.py | 40 +++++++++++++++++++------------------- cyto_dl/nn/vits/utils.py | 2 ++ 4 files changed, 27 insertions(+), 21 deletions(-) diff --git a/cyto_dl/nn/vits/decoder.py b/cyto_dl/nn/vits/decoder.py index 0ef96cda0..fc4883686 100644 --- a/cyto_dl/nn/vits/decoder.py +++ b/cyto_dl/nn/vits/decoder.py @@ -27,6 +27,7 @@ def __init__( emb_dim: Optional[int] = 192, num_layer: Optional[int] = 4, num_head: Optional[int] = 3, + n_output_channels: Optional[int] = 1, has_cls_token: Optional[bool] = False, learnable_pos_embedding: Optional[bool] = True, ) -> None: @@ -66,7 +67,7 @@ def __init__( self.transformer = torch.nn.Sequential( *[Block(emb_dim, num_head) for _ in range(num_layer)] ) - out_dim = torch.prod(torch.as_tensor(patch_size)).item() + out_dim = torch.prod(torch.as_tensor(patch_size)).item() * n_output_channels self.decoder_norm = nn.LayerNorm(emb_dim) self.head = torch.nn.Linear(emb_dim, out_dim) self.num_patches = torch.as_tensor(num_patches) @@ -157,6 +158,7 @@ def __init__( emb_dim: Optional[int] = 192, num_layer: Optional[int] = 4, num_head: Optional[int] = 3, + n_output_channels: Optional[int] = 1, has_cls_token: Optional[bool] = True, learnable_pos_embedding: Optional[bool] = True, ) -> None: @@ -188,6 +190,7 @@ def __init__( emb_dim, num_layer, num_head, + n_output_channels, has_cls_token, learnable_pos_embedding, ) diff --git a/cyto_dl/nn/vits/mae.py b/cyto_dl/nn/vits/mae.py index 17d3e4729..945559d89 100644 --- a/cyto_dl/nn/vits/mae.py +++ b/cyto_dl/nn/vits/mae.py @@ -136,6 +136,7 @@ def __init__( num_layer=decoder_layer, num_head=decoder_head, learnable_pos_embedding=learnable_pos_embedding, + n_output_channels=input_channels, ) @property diff --git a/cyto_dl/nn/vits/seg.py b/cyto_dl/nn/vits/seg.py index 3e10e91a1..3a4591930 100644 --- a/cyto_dl/nn/vits/seg.py +++ b/cyto_dl/nn/vits/seg.py @@ -6,6 +6,7 @@ from monai.networks.blocks import UnetOutBlock, UnetResBlock, UpSample from cyto_dl.nn.vits.mae import MAE_Encoder +from cyto_dl.nn.vits.utils import match_tuple_dimensions class EncodedSkip(torch.nn.Module): @@ -55,7 +56,7 @@ def __init__( self, spatial_dims: int = 3, num_patches: Optional[List[int]] = [2, 32, 32], - base_patch_size: Optional[List[int]] = [4, 8, 8], + patch_size: Optional[List[int]] = [4, 8, 8], emb_dim: Optional[int] = 192, n_decoder_filters: Optional[int] = 16, out_channels: Optional[int] = 6, @@ -63,9 +64,8 @@ def __init__( num_layer: Optional[int] = 3, ) -> None: super().__init__() - total_upsample_factor = np.array(upsample_factor) * np.array(base_patch_size) - self.num_layer = np.min(np.log2(total_upsample_factor)).astype(int) - print(self.num_layer) + total_upsample_factor = np.array(upsample_factor) * np.array(patch_size) + self.num_layer = np.min(np.log2(total_upsample_factor)).astype(int) - 1 residual_resize_factor = list(total_upsample_factor / 2**self.num_layer) input_n_decoder_filters = n_decoder_filters @@ -146,7 +146,15 @@ def forward( features: torch.Tensor, ): # remove global token - features = features[:, 1:] + if len(features.shape) == 4: + assert ( + features.shape[0] == self.num_layer + 1 + ), f"Please set the `n_intermediate_weights` for your encoder to {self.num_layer+1}" + features = features[:, 1:] + elif len(features.shape) == 3: + features = features[1:] + # duplicate features to match num decoder layers + features = torch.stack([features] * (self.num_layer + 1)) prev = None for i in range(self.num_layer + 1): skip = self.upsampling[f"layer_{i}"]["skip"](features[i]) @@ -163,7 +171,7 @@ def __init__( self, spatial_dims: int = 3, num_patches: Optional[List[int]] = [2, 32, 32], - base_patch_size: Optional[List[int]] = [16, 16, 16], + patch_size: Optional[List[int]] = [16, 16, 16], emb_dim: Optional[int] = 768, decoder_layer: Optional[int] = 3, n_decoder_filters: Optional[int] = 16, @@ -180,7 +188,7 @@ def __init__( Number of spatial dimensions num_patches: Optional[List[int]]=[2, 32, 32] Number of patches in each dimension (ZYX) order - base_patch_size: Optional[List[int]]=[16, 16, 16] + patch_size: Optional[List[int]]=[16, 16, 16] Base patch size in each dimension (ZYX) order emb_dim: Optional[int] =768 Embedding dimension of ViT backbone @@ -202,21 +210,13 @@ def __init__( Path to pretrained ViT backbone checkpoint """ super().__init__() - assert spatial_dims in (2, 3) - if isinstance(num_patches, int): - num_patches = [num_patches] * spatial_dims - if isinstance(base_patch_size, int): - base_patch_size = [base_patch_size] * spatial_dims - if isinstance(upsample_factor, int): - upsample_factor = [upsample_factor] * spatial_dims - assert len(num_patches) == spatial_dims - assert len(base_patch_size) == spatial_dims - assert len(upsample_factor) == spatial_dims - + num_patches, patch_size, upsample_factor = match_tuple_dimensions( + spatial_dims, [num_patches, patch_size, upsample_factor] + ) self.encoder = MAE_Encoder( spatial_dims=spatial_dims, num_patches=num_patches, - base_patch_size=base_patch_size, + patch_size=patch_size, emb_dim=emb_dim, **encoder_kwargs, ) @@ -237,7 +237,7 @@ def __init__( self.decoder = SuperresDecoder( spatial_dims=spatial_dims, num_patches=num_patches, - base_patch_size=base_patch_size, + patch_size=patch_size, emb_dim=emb_dim, num_layer=decoder_layer, n_decoder_filters=n_decoder_filters, diff --git a/cyto_dl/nn/vits/utils.py b/cyto_dl/nn/vits/utils.py index aa96557fe..a79910e2b 100644 --- a/cyto_dl/nn/vits/utils.py +++ b/cyto_dl/nn/vits/utils.py @@ -16,7 +16,9 @@ def take_indexes(sequences, indexes): def random_indexes(size: int, device): + # forward indices match original index to shuffled index forward_indexes = torch.randperm(size, device=device, dtype=torch.long) + # backward indices match shuffled index to original index backward_indexes = torch.argsort(forward_indexes) return forward_indexes, backward_indexes