import torch.nn as nn import torch.nn.functional as F import torch import numpy as np import wandb from peft import LoraConfig, get_peft_model from torchvision.transforms import v2 from .unet_base import BaseUnet class SplitQKVLinear(nn.Module): def __init__(self, base_layer: nn.Linear): super().__init__() if base_layer.out_features % 3 != 0: raise ValueError("Expected fused qkv projection with out_features divisible by 3.") self.proj_dim = base_layer.out_features // 3 self.q_proj = nn.Linear(base_layer.in_features, self.proj_dim, bias=base_layer.bias is not None) self.k_proj = nn.Linear(base_layer.in_features, self.proj_dim, bias=base_layer.bias is not None) self.v_proj = nn.Linear(base_layer.in_features, self.proj_dim, bias=base_layer.bias is not None) with torch.no_grad(): self.q_proj.weight.copy_(base_layer.weight[: self.proj_dim]) self.k_proj.weight.copy_(base_layer.weight[self.proj_dim : 2 * self.proj_dim]) self.v_proj.weight.copy_(base_layer.weight[2 * self.proj_dim :]) if base_layer.bias is not None: self.q_proj.bias.copy_(base_layer.bias[: self.proj_dim]) self.k_proj.bias.copy_(base_layer.bias[self.proj_dim : 2 * self.proj_dim]) self.v_proj.bias.copy_(base_layer.bias[2 * self.proj_dim :]) def forward(self, x: torch.Tensor) -> torch.Tensor: q = self.q_proj(x) k = self.k_proj(x) v = self.v_proj(x) return torch.cat((q, k, v), dim=-1) def replace_qkv_with_split_projections(module: nn.Module): replaced_layers = 0 for child_name, child_module in list(module.named_children()): if ( child_name == "qkv" and isinstance(child_module, nn.Linear) and child_module.out_features == child_module.in_features * 3 ): setattr(module, child_name, SplitQKVLinear(child_module)) replaced_layers += 1 continue replaced_layers += replace_qkv_with_split_projections(child_module) return replaced_layers def apply_lora_to_qv_projections(module: nn.Module, rank: int, alpha: float | None = None): replaced_layers = replace_qkv_with_split_projections(module) if replaced_layers == 0: raise ValueError("No qkv projections found in MedSAM image encoder for LoRA injection.") lora_config = LoraConfig( r=rank, lora_alpha=(alpha if alpha is not None else rank), target_modules=["q_proj", "v_proj"], bias="none", ) peft_module = get_peft_model(module, lora_config) return peft_module, replaced_layers def pad_to_2d(x: torch.Tensor, stride: int): h, w = x.shape[-2:] new_h = h if h % stride == 0 else h + stride - (h % stride) new_w = w if w % stride == 0 else w + stride - (w % stride) top = (new_h - h) // 2 bottom = (new_h - h) - top left = (new_w - w) // 2 right = (new_w - w) - left pads = (left, right, top, bottom) x_pad = F.pad(x, pads, mode="constant", value=0) return x_pad, pads def unpad_2d(x: torch.Tensor, pads): left, right, top, bottom = pads if top or bottom: end_h = -bottom if bottom > 0 else None x = x[:, :, top:end_h, :] if left or right: end_w = -right if right > 0 else None x = x[:, :, :, left:end_w] return x class ConvBlock(nn.Module): def __init__( self, in_channels, out_channels, conv_kwargs={"kernel_size": 3, "stride": 1, "padding": 1}, ): super(ConvBlock, self).__init__() self.block = nn.Sequential( nn.Conv2d(in_channels, out_channels, **conv_kwargs), nn.InstanceNorm2d(out_channels), nn.LeakyReLU(), ) def forward(self, x): return self.block(x) # Single encoder block class DownConvBlock(nn.Module): def __init__( self, in_channels: list, out_channels: list, conv_kwargs={"kernel_size": 3, "stride": 1, "padding": 1}, ): super(DownConvBlock, self).__init__() assert len(in_channels) == len( out_channels ), f"in_channels length is {len(in_channels)} while out_channels is {len(out_channels)}" # Variable number of convolutional block in each layer, based on the in_channels and out_channels length self.conv_blocks = nn.ModuleList( [ ConvBlock(in_ch, out_ch, conv_kwargs) for in_ch, out_ch in zip(in_channels, out_channels) ] ) self.pool = nn.MaxPool2d(2, 2) def forward(self, x): for block in self.conv_blocks: x = block(x) return self.pool(x), x # Single decoder block class UpConvBlock(nn.Module): def __init__( self, in_channels: list, out_channels: list, up_conv=True, conv_kwargs={"kernel_size": 3, "stride": 1, "padding": 1}, upconv_kwargs={"kernel_size": 2, "stride": 2}, ): super(UpConvBlock, self).__init__() assert len(in_channels) == len( out_channels ), f"in_channels length is {len(in_channels)} while out_channels is {len(out_channels)}" # Variable number of convolutional block in each layer, based on the in_channels and out_channels length self.conv_blocks = nn.ModuleList( [ ConvBlock(in_ch, out_ch, conv_kwargs) for in_ch, out_ch in zip(in_channels, out_channels) ] ) self.up_conv = up_conv if self.up_conv: self.up_conv_op = nn.ConvTranspose2d( out_channels[-1], out_channels[-1], **upconv_kwargs ) def forward(self, x): for block in self.conv_blocks: x = block(x) if self.up_conv: return self.up_conv_op(x) else: return x class FiLM2d(nn.Module): """ Feature-wise Linear Modulation for a 2-D feature map. (γ, β) are generated from a learned embedding of organ_id. """ def __init__( self, n_organs: int, in_channels: int, emb_dim: int | None = None, hidden: int | None = None, ): super().__init__() hidden = hidden or 2 * in_channels self.embed = nn.Embedding(n_organs + 1, emb_dim) self.mlp = nn.Sequential( nn.Linear(emb_dim, hidden), nn.ReLU(inplace=True), nn.Linear(hidden, 2 * in_channels), # → [β‖γ] ) # initialise so that FiLM starts as identity: γ≈1, β≈0 nn.init.zeros_(self.mlp[-1].weight) nn.init.constant_(self.mlp[-1].bias[:in_channels], 0) # β nn.init.constant_(self.mlp[-1].bias[in_channels:], 1) # γ def compute_modulation(self, organ_id: torch.Tensor): """Return FiLM coefficients and conditioning embedding for a batch.""" B = organ_id.shape[0] device = organ_id.device mask = organ_id >= 0 # [B] q_org = self.embed(organ_id.clamp(min=0)) # [B, D] (dummy for unknown) q_img = self.embed( torch.tensor([self.embed.weight.shape[0] - 1], device=device).expand(B) ) q = torch.where(mask[:, None], q_org, q_img) beta_gamma = self.mlp(q) # (B, 2C) beta, gamma = beta_gamma.chunk(2, dim=-1) # each (B, C) beta = beta.unsqueeze(-1).unsqueeze(-1) gamma = gamma.unsqueeze(-1).unsqueeze(-1) return gamma, beta, q def forward(self, x: torch.Tensor, organ_id: torch.Tensor): """ x : (B, C, H, W) organ_id : (B,) integer 0…n_organs-1 """ gamma, beta, _ = self.compute_modulation(organ_id) return gamma * x + beta class SharedFiLMModulator: """ Compatibility helper that exposes the same collection interface used by UNet2DAttn notebooks: compute all per-layer FiLM coefficients in one call. """ def __init__(self, model: "UNet2DFiLM"): self.model = model def compute_all_modulations( self, original_img: torch.Tensor, organ_id: torch.Tensor, layer_configs: list[tuple[int, int]], return_compact: bool = False, ): del original_img # FiLM modulation depends only on organ_id in this model. modulations = {} compact_modulations = {} if return_compact else None film_layers = self.model._get_film_layers_in_order() if len(film_layers) != len(layer_configs): raise RuntimeError( f"FiLM layer/config mismatch: {len(film_layers)} layers vs " f"{len(layer_configs)} configs." ) for (layer_id, n_channels), film_layer in zip(layer_configs, film_layers): gamma, beta, film_context = film_layer.compute_modulation(organ_id) modulations[layer_id] = (gamma, beta) if not return_compact: continue gamma_flat = gamma.squeeze(-1).squeeze(-1) beta_flat = beta.squeeze(-1).squeeze(-1) if n_channels != self.model.max_channels: gamma_compact = F.interpolate( gamma_flat.unsqueeze(1), size=self.model.max_channels, mode="linear", align_corners=False, ).squeeze(1) beta_compact = F.interpolate( beta_flat.unsqueeze(1), size=self.model.max_channels, mode="linear", align_corners=False, ).squeeze(1) else: gamma_compact = gamma_flat beta_compact = beta_flat compact_modulations[layer_id] = ( gamma_compact, beta_compact, film_context, ) projected_shapes = torch.empty(0, device=organ_id.device) if return_compact: return modulations, projected_shapes, compact_modulations return modulations, projected_shapes class DownConvBlockFiLM(nn.Module): """ Conv → FiLM → Conv → FiLM → Pool. Except for FiLM, the API and behaviour remain the same. """ def __init__( self, in_channels: list[int], out_channels: list[int], n_organs: int, conv_kwargs={"kernel_size": 3, "stride": 1, "padding": 1}, emb_dim: int = 64, ): super().__init__() assert len(in_channels) == len( out_channels ), f"in_channels length is {len(in_channels)} while out_channels is {len(out_channels)}" self.conv_blocks = nn.ModuleList( [ ConvBlock(in_ch, out_ch, conv_kwargs) for in_ch, out_ch in zip(in_channels, out_channels) ] ) self.film_blocks = nn.ModuleList( [ FiLM2d( n_organs=n_organs, in_channels=out_ch, emb_dim=emb_dim, ) for out_ch in out_channels ] ) self.pool = nn.MaxPool2d(2, 2) def forward( self, x: torch.Tensor, organ_id: torch.Tensor | None = None, gammas: list[torch.Tensor] | None = None, betas: list[torch.Tensor] | None = None, ): """ organ_id : (B,) torch.long (0 = breast, 1 = thyroid, …) """ if (gammas is None) != (betas is None): raise ValueError("gammas and betas must be provided together.") for idx, (conv, film) in enumerate(zip(self.conv_blocks, self.film_blocks)): x = conv(x) # usual conv-norm-ReLU if gammas is None: if organ_id is None: raise ValueError("organ_id is required when FiLM modulations are not precomputed.") x = film(x, organ_id) # FiLM modulation else: x = gammas[idx] * x + betas[idx] return self.pool(x), x # (downsampled, skip-connection) class UpConvBlockFiLM(nn.Module): """ Up-sampling block with FiLM conditioning. Conv → FiLM → Conv → FiLM → (optional) ConvTranspose2d """ def __init__( self, in_channels: list[int], out_channels: list[int], n_organs: int, up_conv: bool = True, conv_kwargs: dict = {"kernel_size": 3, "stride": 1, "padding": 1}, upconv_kwargs: dict = {"kernel_size": 2, "stride": 2}, emb_dim: int = 64, ): super().__init__() assert len(in_channels) == len( out_channels ), f"in_channels length is {len(in_channels)} while out_channels is {len(out_channels)}" self.conv_blocks = nn.ModuleList( [ ConvBlock(in_ch, out_ch, conv_kwargs) for in_ch, out_ch in zip(in_channels, out_channels) ] ) self.film_blocks = nn.ModuleList( [ FiLM2d( n_organs=n_organs, in_channels=out_ch, emb_dim=emb_dim, ) for out_ch in out_channels ] ) self.up_conv = up_conv if self.up_conv: self.up_conv_op = nn.ConvTranspose2d( out_channels[-1], out_channels[-1], **upconv_kwargs ) def forward( self, x: torch.Tensor, organ_id: torch.Tensor | None = None, gammas: list[torch.Tensor] | None = None, betas: list[torch.Tensor] | None = None, ): """ x : (B, C, H, W) organ_id : (B,) long tensor - 0=breast, 1=thyroid, … """ if (gammas is None) != (betas is None): raise ValueError("gammas and betas must be provided together.") for idx, (conv, film) in enumerate(zip(self.conv_blocks, self.film_blocks)): x = conv(x) if gammas is None: if organ_id is None: raise ValueError("organ_id is required when FiLM modulations are not precomputed.") x = film(x, organ_id) else: x = gammas[idx] * x + betas[idx] if self.up_conv: x = self.up_conv_op(x) return x class UNet2DFiLM(BaseUnet): def __init__( self, in_channels: int, num_classes: int, n_organs: int, size: int = 32, depth: int = 3, *, film_start: int = 0, use_film: bool = True, film_embed: int = 64, distill: bool = False, distill_unet: bool = False, medsam_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/medsam_unfreezed/model.safetensors", unet_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/unet5_attn/model.safetensors", unet_teacher_kwargs: dict | None = None, ): """ UNet with symmetric FiLM conditioning in encoder and decoder. """ super().__init__( in_channels=in_channels, num_classes=num_classes, n_organs=n_organs, size=size, depth=depth, film_start=film_start, use_film=use_film, film_embed=film_embed, distill=distill, distill_unet=distill_unet, medsam_teacher_ckpt=medsam_teacher_ckpt, unet_teacher_ckpt=unet_teacher_ckpt, unet_teacher_kwargs=unet_teacher_kwargs, ) def _build_model( self, *, film_start: int = 0, use_film: bool = True, film_embed: int = 64, distill: bool = False, distill_unet: bool = False, medsam_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/medsam_unfreezed/model.safetensors", unet_teacher_ckpt: str = "/work/phd_ultrasounds/UUSIC_new/checkpoints/unet5_attn/model.safetensors", unet_teacher_kwargs: dict | None = None, **kwargs, ): if kwargs: unknown = ", ".join(sorted(kwargs.keys())) raise TypeError(f"Unexpected UNet2DFiLM kwargs: {unknown}") self.film_start = max(0, int(film_start)) self.use_film = bool(use_film) self.use_attn = self.use_film self.attn_start = self.film_start self.film_embed = int(film_embed) self.max_channels = self.size * (2 ** (self.depth + 1)) self.criterion = DiceBCELoss() self.shared_attn = SharedFiLMModulator(self) if self.use_film else None # ---------------- Encoder ---------------- self.encoder = nn.ModuleDict() if self.use_film and 0 >= self.film_start: self.encoder["0"] = DownConvBlockFiLM( [self.in_channels, self.size], [self.size, self.size * 2], n_organs=self.n_organs, emb_dim=self.film_embed, ) else: self.encoder["0"] = DownConvBlock( [self.in_channels, self.size], [self.size, self.size * 2] ) for i in range(1, self.depth): in_ch = [self.size * (2**i), self.size * (2**i)] out_ch = [self.size * (2**i), self.size * (2 ** (i + 1))] key = str(i) if self.use_film and i >= self.film_start: self.encoder[key] = DownConvBlockFiLM( in_ch, out_ch, n_organs=self.n_organs, emb_dim=self.film_embed, ) else: self.encoder[key] = DownConvBlock(in_ch, out_ch) # ---------------- Bottleneck ---------------- if self.use_film: self.bottleneck = UpConvBlockFiLM( [self.size * (2**self.depth), self.size * (2**self.depth)], [self.size * (2**self.depth), self.size * (2 ** (self.depth + 1))], n_organs=self.n_organs, emb_dim=self.film_embed, ) else: self.bottleneck = UpConvBlock( [self.size * (2**self.depth), self.size * (2**self.depth)], [self.size * (2**self.depth), self.size * (2 ** (self.depth + 1))], ) # ---------------- Decoder ---------------- self.decoder = nn.ModuleDict() for i in range(self.depth, 1, -1): use_film_at_level = self.use_film and (i - 1) >= self.film_start if use_film_at_level: self.decoder[str(i - 1)] = UpConvBlockFiLM( [ self.size * (2 ** (i + 1)) + self.size * (2**i), self.size * (2**i), ], [self.size * (2**i), self.size * (2**i)], n_organs=self.n_organs, emb_dim=self.film_embed, ) else: self.decoder[str(i - 1)] = UpConvBlock( [ self.size * (2 ** (i + 1)) + self.size * (2**i), self.size * (2**i), ], [self.size * (2**i), self.size * (2**i)], ) if self.use_film and 0 >= self.film_start: self.decoder["0"] = UpConvBlockFiLM( [self.size * 4 + self.size * 2, self.size * 2], [self.size * 2, self.size * 2], n_organs=self.n_organs, up_conv=False, emb_dim=self.film_embed, ) else: self.decoder["0"] = UpConvBlock( [self.size * 4 + self.size * 2, self.size * 2], [self.size * 2, self.size * 2], up_conv=False, ) self.out_layer = ConvBlock( self.size * 2, self.out_channels, conv_kwargs={"kernel_size": 1, "stride": 1, "padding": 0}, ) self._init_distillation( distill=distill, distill_unet=distill_unet, medsam_teacher_ckpt=medsam_teacher_ckpt, unet_teacher_ckpt=unet_teacher_ckpt, unet_teacher_kwargs=unet_teacher_kwargs, ) def _build_layer_configs(self): """Match UNet2DAttn layer ordering for notebook-side FiLM collection.""" configs = [] layer_id = 0 for i in range(self.depth): if self.use_film and i >= self.film_start: configs.append((layer_id, self.size * (2**i))) layer_id += 1 configs.append((layer_id, self.size * (2 ** (i + 1)))) layer_id += 1 if self.use_film: configs.append((layer_id, self.size * (2**self.depth))) layer_id += 1 configs.append((layer_id, self.size * (2 ** (self.depth + 1)))) layer_id += 1 for i in range(self.depth - 1, -1, -1): if self.use_film and i >= self.film_start: configs.append((layer_id, self.size * (2 ** (i + 1)))) layer_id += 1 configs.append((layer_id, self.size * (2 ** (i + 1)))) layer_id += 1 return configs def _get_film_layers_in_order(self): """Return FiLM modules in the same order as `_build_layer_configs`.""" film_layers = [] for i in range(self.depth): layer = self.encoder[str(i)] if isinstance(layer, DownConvBlockFiLM): film_layers.extend(layer.film_blocks) if isinstance(self.bottleneck, UpConvBlockFiLM): film_layers.extend(self.bottleneck.film_blocks) for i in range(self.depth - 1, -1, -1): layer = self.decoder[str(i)] if isinstance(layer, UpConvBlockFiLM): film_layers.extend(layer.film_blocks) return film_layers def _prepare_forward( self, *, pixel_values: torch.Tensor, organ_id: torch.Tensor | None = None, **kwargs, ) -> dict: forward_ctx = { "mod_list": None, "mod_idx": 0, } if self.use_film: if organ_id is None: raise ValueError("organ_id is required when FiLM is enabled.") layer_configs = self._build_layer_configs() film_layers = self._get_film_layers_in_order() if len(film_layers) != len(layer_configs): raise RuntimeError( f"FiLM layer/config mismatch: {len(film_layers)} layers vs " f"{len(layer_configs)} configs." ) modulations = { layer_id: film_layer.compute_modulation(organ_id)[:2] for (layer_id, _), film_layer in zip(layer_configs, film_layers) } forward_ctx["mod_list"] = [modulations[i] for i in range(len(layer_configs))] return forward_ctx def _next_modulation(self, forward_ctx: dict): mod_idx = forward_ctx["mod_idx"] mod_list = forward_ctx["mod_list"] gammas = [mod_list[mod_idx][0], mod_list[mod_idx + 1][0]] betas = [mod_list[mod_idx][1], mod_list[mod_idx + 1][1]] forward_ctx["mod_idx"] = mod_idx + 2 return gammas, betas def _encode( self, layer, x, organ_id=None, forward_ctx=None, ): if isinstance(layer, DownConvBlockFiLM): if forward_ctx is not None and forward_ctx.get("mod_list") is not None: gammas, betas = self._next_modulation(forward_ctx) return layer(x, gammas=gammas, betas=betas) return layer(x, organ_id=organ_id) else: return layer(x) def _decode( self, layer, x, organ_id=None, forward_ctx=None, ): if isinstance(layer, UpConvBlockFiLM): if forward_ctx is not None and forward_ctx.get("mod_list") is not None: gammas, betas = self._next_modulation(forward_ctx) return layer(x, gammas=gammas, betas=betas) return layer(x, organ_id=organ_id) else: return layer(x) def _bottleneck( self, x, organ_id=None, forward_ctx=None, ): if isinstance(self.bottleneck, UpConvBlockFiLM): if forward_ctx is not None and forward_ctx.get("mod_list") is not None: gammas, betas = self._next_modulation(forward_ctx) return self.bottleneck(x, gammas=gammas, betas=betas) return self.bottleneck(x, organ_id=organ_id) else: return self.bottleneck(x) def __str__(self): model_parameters = filter(lambda p: p.requires_grad, self.parameters()) params = sum([np.prod(p.size()) for p in model_parameters]) film_status = "enabled" if self.use_film else "disabled" film_range = f"from level {self.film_start}" if self.use_film else "N/A" return ( super().__str__() + f"\nTrainable parameters: {params}" f"\nFiLM: {film_status} ({film_range})" ) class DiceBCELoss(nn.Module): def __init__(self, dice_weight: float = 1.0, bce_weight: float = 1.0): super().__init__() self.dice_weight = dice_weight self.bce_weight = bce_weight self.eps = 1e-6 def forward(self, logits: torch.Tensor, gt: torch.Tensor) -> torch.Tensor: gt = gt.float() bce = F.binary_cross_entropy_with_logits( logits.squeeze(), gt.squeeze(), reduction="mean" ) # Soft Dice loss probs = torch.sigmoid(logits) dims = tuple(range(2, probs.dim())) # (H, W) or (D,H,W) # per‑class Dice, per‑sample inter = (probs * gt).sum(dims) * 2 union = probs.sum(dims) + gt.sum(dims) dice = 1 - (inter + self.eps) / (union + self.eps) # [B, C] dice = dice.mean() loss = self.dice_weight * dice + self.bce_weight * bce return loss class MedSAM(nn.Module): def __init__( self, image_encoder, mask_decoder, prompt_encoder, freeze_image_encoder=True, predict_bboxes=False, image_encoder_lora_rank=0, ): super().__init__() self.image_encoder = image_encoder self.mask_decoder = mask_decoder self.prompt_encoder = prompt_encoder self.criterion = DiceBCELoss() self.freeze_image_encoder = freeze_image_encoder self.predict_bboxes = predict_bboxes if self.freeze_image_encoder: for param in self.image_encoder.parameters(): param.requires_grad = False self.image_encoder_lora_rank = int(image_encoder_lora_rank) self.image_encoder_lora_layers = 0 if self.image_encoder_lora_rank > 0: self.image_encoder, self.image_encoder_lora_layers = apply_lora_to_qv_projections( self.image_encoder, rank=self.image_encoder_lora_rank, ) # Classification head self.multi_cls = nn.Sequential(nn.Linear(256 * 64 * 64, 10)) # Bounding box regression head self.bbox_regr = nn.Sequential(nn.Linear(256 * 64 * 64, 4)) # Learnable prompt embeddings (no input required) self.learned_sparse_embeddings = nn.Parameter( torch.randn(1, 2, 256) # (batch, num_tokens, embed_dim) ) self.learned_dense_embeddings = nn.Parameter( torch.randn(1, 256, 64, 64) # (batch, embed_dim, H, W) ) def forward( self, pixel_values, organ_id=None, labels=None, masks=None, bbox_coords=None, organ_id_metric=None, pixel_values_medsam=None, **kwargs, # ignored, for peft compatibility ): batch_size = pixel_values.shape[0] # Get image embeddings image_embedding = self.image_encoder(pixel_values_medsam) # (B, 256, 64, 64) # Classification output emb_flattened = torch.flatten(image_embedding, 1) multi_cls_out = self.multi_cls(emb_flattened) # Bounding box output (if enabled) if self.predict_bboxes: bbox_out = self.bbox_regr(emb_flattened) else: bbox_out = None # Expand learned embeddings to batch size sparse_embeddings = self.learned_sparse_embeddings dense_embeddings = self.learned_dense_embeddings # Get positional encoding image_pe = self.prompt_encoder.get_dense_pe() # (1, 256, 64, 64) # Decode mask low_res_masks, iou_predictions = self.mask_decoder( image_embeddings=image_embedding, # (B, 256, 64, 64) image_pe=image_pe, # (1, 256, 64, 64) sparse_prompt_embeddings=sparse_embeddings, # (B, 2, 256) dense_prompt_embeddings=dense_embeddings, # (B, 256, 64, 64) multimask_output=False, ) # (B, 1, 256, 256) if masks is not None: loss = self.criterion( low_res_masks.squeeze(1), v2.functional.resize(masks, (256, 256)) ) else: loss = 0.0 return { "loss": loss, "logits": low_res_masks.squeeze(1), "labels": masks, "organ_id": organ_id, "organ_id_metric": organ_id_metric, } class DistillationLoss(nn.Module): """ Improved distillation loss with multiple components: - Cosine similarity loss (direction alignment) - MSE loss (magnitude alignment) - Optional L1 loss (sparsity) """ def __init__(self, temperature=3.0, alpha=0.5, use_l1=False): super().__init__() self.temperature = temperature self.alpha = alpha # Weight between cosine and MSE self.use_l1 = use_l1 def forward(self, student_logits, teacher_logits, tau=0.7): """ Args: student_logits: (B, D) - student features teacher_logits: (B, D) - teacher features """ # # Normalize features for stable training # student_norm = F.normalize(student_logits, p=2, dim=1) # teacher_norm = F.normalize(teacher_logits, p=2, dim=1) # # Cosine similarity loss (encourages directional alignment) # cosine_loss = ( # 1 - F.cosine_similarity(student_logits, teacher_logits, dim=1).mean() # ) # # MSE loss on normalized features (encourages magnitude alignment) # mse_loss = F.mse_loss(student_norm, teacher_norm) # # Combined loss # loss = self.alpha * cosine_loss + (1 - self.alpha) * mse_loss # # Optional L1 for sparsity # if self.use_l1: # l1_loss = F.l1_loss(student_logits, teacher_logits) # loss = loss + 0.1 * l1_loss B = student_logits.size(0) s = student_logits.flatten(start_dim=1) t = teacher_logits.flatten(start_dim=1).detach() s = F.normalize(s, dim=-1) t = F.normalize(t, dim=-1) d2 = (s[:, None, :] - t[None, :, :]).pow(2).sum(dim=-1) logits = -d2 / tau labels = torch.arange(B, device=s.device) loss = F.cross_entropy(logits, labels) return { "loss": loss, } class MedSAMPrompt(nn.Module): def __init__( self, image_encoder, mask_decoder, prompt_encoder, freeze_image_encoder=True, predict_bboxes=False, n_organs=1, image_encoder_lora_rank=0, ): super().__init__() self.image_encoder = image_encoder self.mask_decoder = mask_decoder self.prompt_encoder = prompt_encoder self.criterion = DiceBCELoss() self.freeze_image_encoder = freeze_image_encoder self.predict_bboxes = predict_bboxes if self.freeze_image_encoder: for param in self.image_encoder.parameters(): param.requires_grad = False self.image_encoder_lora_rank = int(image_encoder_lora_rank) self.image_encoder_lora_layers = 0 if self.image_encoder_lora_rank > 0: self.image_encoder, self.image_encoder_lora_layers = apply_lora_to_qv_projections( self.image_encoder, rank=self.image_encoder_lora_rank, ) # Classification head self.multi_cls = nn.Sequential(nn.Linear(256 * 64 * 64, 10)) # Bounding box regression head self.bbox_regr = nn.Sequential(nn.Linear(256 * 64 * 64, 4)) # Learnable prompt embeddings (no input required) self.sparse_embeddings = nn.Parameter(torch.randn(1, 2, 256)) self.dense_embeddings = nn.Embedding(n_organs + 1, 1 * 256 * 64 * 64) # self.sparse_embeddings = nn.ParameterDict( # { # str(id_): nn.Parameter(torch.randn(1, 2, 256)) # for id_ in set(organ_to_class_dict.values()) # } # ) # self.dense_embeddings = nn.ParameterDict( # { # str(id_): nn.Parameter(torch.randn(1, 256, 64, 64)) # for id_ in set(organ_to_class_dict.values()) # } # ) def forward( self, pixel_values, organ_id=None, labels=None, masks=None, bbox_coords=None, organ_id_metric=None, pixel_values_medsam=None, **kwargs, # ignored, for peft compatibility ): batch_size = pixel_values.shape[0] B = pixel_values.shape[0] # Get image embeddings image_embedding = self.image_encoder(pixel_values_medsam) # (B, 256, 64, 64) # Classification output emb_flattened = torch.flatten(image_embedding, 1) multi_cls_out = self.multi_cls(emb_flattened) # Bounding box output (if enabled) if self.predict_bboxes: bbox_out = self.bbox_regr(emb_flattened) else: bbox_out = None if organ_id is None: raise ValueError("organ_id must be provided for selecting embeddings.") sparse_emb_list, dense_emb_list = [], [] mask = organ_id >= 0 # [B] idx = torch.where(mask, organ_id, self.dense_embeddings.weight.shape[0] - 1) dense_embeddings = self.dense_embeddings(idx).view(B, 256, 64, 64) # Get positional encoding image_pe = self.prompt_encoder.get_dense_pe() # (1, 256, 64, 64) # Decode mask low_res_masks, iou_predictions = self.mask_decoder( image_embeddings=image_embedding, # (B, 256, 64, 64) image_pe=image_pe, # (1, 256, 64, 64) sparse_prompt_embeddings=self.sparse_embeddings, # (B, 2, 256) dense_prompt_embeddings=dense_embeddings, # (B, 256, 64, 64) multimask_output=False, ) # (B, 1, 256, 256) if masks is not None: loss = self.criterion( low_res_masks.squeeze(1), v2.functional.resize(masks, (256, 256)) ) else: loss = 0.0 return { "loss": loss, "logits": low_res_masks.squeeze(1), "labels": masks, "organ_id": organ_id, "organ_id_metric": organ_id_metric, } # config = {"in_channels": 3,"num_classes": 1,"n_organs": 8,"size": 32,"depth": 5,"film_start": 0,"use_film": 1}