""" LTF V3.3: V3 + pool-logit gated fusion. Motivation: - Keep V3 outputs (logits_level1 / logits_level2 / logits_fusion) unchanged for training tools. - Use "logit_pool" signals (attention-weighted instance logits) as a reliability cue to gate fusion between level1_emb and level2_emb, while keeping fusion_head input dim = 2*embed_dim. """ from typing import Optional, Tuple import torch import torch.nn as nn from model.LTF import LTFCore from model.LTF_v2 import ( BagTokenShuffle, _clip_vit_resblock_modules, _count_clip_vit_depth, ) from mil_module.mil_core import MILCore class LTFCoreV3_3(nn.Module): """ V3 core + a tiny gate that uses logit_pool(level1_tiles) and logit_pool(level2_cls-bag). - Level1: patch tokens (tiles only) -> MIL per tile -> tile bag embeddings -> query attention -> level1_emb Additionally: compute per-tile instance logits and attention weights -> logit_pool across tiles. - Level2: CLS tokens (global + tiles) -> MIL (attn pooling) -> bag_embed = level2_emb Additionally: compute logit_pool from MILHead's attn * inst_logits. - Fusion: g = sigmoid(MLP([l1_pool_logits, l2_pool_logits])) in (0,1), fusion_in = [g*level1_emb, (1-g)*level2_emb] -> fusion_head (same as V3 input dim). """ def __init__( self, embed_dim=1024, num_heads=8, num_classes=2, num_layers=1, dropout=0.1, tile_heads=2, tile_layers=2, use_global_image=True, use_patch_transformer=True, linear_classifier=False, mil_heads=8, fusion_hidden=512, l1_cross_shuffle_p: float = 0.9, l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9), l1_with_replacement_p: float = 0.5, l2_cross_shuffle_p: float = 0.5, l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5), l2_with_replacement_p: float = 0.5, bag_token_shuffle_p: float = 0.5, cross_shuffle_min_keep_original: int = 0, use_mil_cross_shuffle: bool = True, pool_gate_hidden: int = 16, use_pool_gate_fusion: bool = True, ): super().__init__() self.embed_dim = embed_dim self.use_global_image = use_global_image self.num_classes = num_classes self.use_mil_cross_shuffle = bool(use_mil_cross_shuffle) self.l1_cross_shuffle_p = float(l1_cross_shuffle_p) self.l1_replace_ratio_range = ( float(l1_replace_ratio_range[0]), float(l1_replace_ratio_range[1]), ) self.l1_with_replacement_p = float(l1_with_replacement_p) self.l2_cross_shuffle_p = float(l2_cross_shuffle_p) self.l2_replace_ratio_range = ( float(l2_replace_ratio_range[0]), float(l2_replace_ratio_range[1]), ) self.l2_with_replacement_p = float(l2_with_replacement_p) self.bag_token_shuffle_p = float(bag_token_shuffle_p) self.cross_shuffle_min_keep_original = int(cross_shuffle_min_keep_original) self.use_pool_gate_fusion = bool(use_pool_gate_fusion) # Keep the same core structure and frozen baseline classifier as V3 for compatibility. self.baseline_core = LTFCore( embed_dim=embed_dim, num_heads=num_heads, num_classes=num_classes, num_layers=num_layers, dropout=dropout, tile_heads=tile_heads, tile_layers=tile_layers, use_global_image=use_global_image, use_patch_transformer=use_patch_transformer, linear_classifier=linear_classifier, ) for p in self.baseline_core.classifier.parameters(): p.requires_grad = False mil_kw = dict( attn_hidden=min(256, embed_dim), attn_dropout=dropout, attn_bag_mode="embed_cls", ) self.mil_level1 = MILCore( d_in=embed_dim, num_classes=num_classes, mil_head_pooling="attn", mil_head_kwargs=mil_kw, use_cross_shuffle=self.use_mil_cross_shuffle, trans=BagTokenShuffle(), allow_trans_token_drop=False, ) self.mil_level2 = MILCore( d_in=embed_dim, num_classes=num_classes, mil_head_pooling="attn", mil_head_kwargs=mil_kw, use_cross_shuffle=self.use_mil_cross_shuffle, trans=BagTokenShuffle(), allow_trans_token_drop=False, ) # Level-1 aggregation across tiles (query attention) + tile-level instance logits for logit_pool. self.level1_query = nn.Parameter(torch.zeros(1, 1, embed_dim)) mh = mil_heads if embed_dim % mil_heads == 0 else 8 while embed_dim % mh != 0: mh -= 1 if mh < 1: mh = 1 self.level1_aggregate_attn = nn.MultiheadAttention( embed_dim, mh, dropout=dropout, batch_first=True ) # Heads self.head_level1 = nn.Linear(embed_dim, num_classes) self.head_level2 = nn.Linear(embed_dim, num_classes) self.tile_inst_head = nn.Linear(embed_dim, num_classes) fusion_in = embed_dim + embed_dim self.fusion_head = nn.Sequential( nn.Linear(fusion_in, fusion_hidden), nn.ReLU(), nn.Dropout(p=dropout), nn.Linear(fusion_hidden, num_classes), ) # Gate uses pooled logits from L1/L2: concat => [B, 2*C] gate_in = 2 * num_classes self.pool_gate = nn.Sequential( nn.Linear(gate_in, int(pool_gate_hidden)), nn.ReLU(), nn.Linear(int(pool_gate_hidden), 1), ) nn.init.trunc_normal_(self.level1_query, std=0.02) @staticmethod def _sample_replace_ratio(low: float, high: float, device: torch.device) -> float: lo, hi = min(low, high), max(low, high) lo = max(0.0, lo) hi = min(1.0, hi) if hi <= lo: return float(lo) u = torch.rand((), device=device) return float(lo + (hi - lo) * u) def _apply_patch_transformer(self, x): b2, n2, t2, d2 = x.shape xf = x.reshape(b2 * n2, t2, d2) if self.baseline_core.patch_transformer is not None: xf = self.baseline_core.patch_transformer(xf) return xf.reshape(b2, n2, t2, d2) @staticmethod def _logit_pool_from_mil_head(mil_head_out: dict) -> torch.Tensor: """ Compute logit_pool = sum(attn * inst_logits) over instances. Requires pooling='attn' so that 'attn' is not None. """ inst = mil_head_out.get("inst_logits") attn = mil_head_out.get("attn") if inst is None or attn is None: raise RuntimeError("logit_pool requires MILHead outputs: inst_logits and attn.") return (attn * inst).sum(dim=1) def forward( self, x, labels: Optional[torch.Tensor] = None, x_mil: Optional[torch.Tensor] = None, return_viz: bool = False, ): x_feat = self._apply_patch_transformer(x) if x_mil is not None: if x_mil.shape != x.shape: raise ValueError(f"x_mil shape {tuple(x_mil.shape)} != x shape {tuple(x.shape)}") x_mil_feat = self._apply_patch_transformer(x_mil) else: x_mil_feat = x_feat b, n, t, d = x_mil_feat.shape if n < 2: raise ValueError(f"LTFCoreV3_3 expects N>=2 (global + tiles), got N={n}") nt = n - 1 train_shuffle = self.training and labels is not None p_bag_shuffle = self.bag_token_shuffle_p if train_shuffle else 0.0 # --- Level 1: patch tokens (tiles only, no CLS) --- patch_only = x_mil_feat[:, 1:, 1:, :].reshape(b * nt, t - 1, d) do1 = train_shuffle and float(torch.rand((), device=x.device)) < self.l1_cross_shuffle_p rr1 = self._sample_replace_ratio(*self.l1_replace_ratio_range, x.device) wr1 = bool(float(torch.rand((), device=x.device)) < self.l1_with_replacement_p) y1 = labels.unsqueeze(1).expand(-1, nt).reshape(-1) if labels is not None else None mk = self.cross_shuffle_min_keep_original mil1 = self.mil_level1( patch_only, y=y1, return_dict=True, do_cross_shuffle=do1, replace_ratio=rr1, with_replacement=wr1, detach_inserted=True, cross_shuffle_min_keep=mk, p_shuffle=p_bag_shuffle, ) be = mil1["mil_head"]["bag_embed"] if be is None: raise RuntimeError("MIL level1 expected bag_embed (use attn + embed_cls).") bag_seq = be.reshape(b, nt, d) # per-tile embeddings # Query-attn aggregate -> level1_emb q = self.level1_query.expand(b, -1, -1) agg_out, attn_w = self.level1_aggregate_attn( q, bag_seq, bag_seq, need_weights=True, average_attn_weights=False ) level1_emb = agg_out.squeeze(1) # [B, D] # logit_pool across tiles using the same attention weights. # attn_w: [B, heads, 1, Nt] when average_attn_weights=False if attn_w is None: raise RuntimeError("Expected attention weights from level1_aggregate_attn.") tile_attn = attn_w.mean(dim=1).squeeze(1) # [B, Nt] tile_inst_logits = self.tile_inst_head(bag_seq) # [B, Nt, C] level1_logit_pool = (tile_attn.unsqueeze(-1) * tile_inst_logits).sum(dim=1) # [B, C] # --- Level 2: CLS tokens (global + tiles) --- cls_seq = x_mil_feat[:, :, 0, :] # [B, N, D] do2 = train_shuffle and float(torch.rand((), device=x.device)) < self.l2_cross_shuffle_p rr2 = self._sample_replace_ratio(*self.l2_replace_ratio_range, x.device) wr2 = bool(float(torch.rand((), device=x.device)) < self.l2_with_replacement_p) mil2 = self.mil_level2( cls_seq, y=labels, return_dict=True, do_cross_shuffle=do2, replace_ratio=rr2, with_replacement=wr2, detach_inserted=True, cross_shuffle_min_keep=mk, p_shuffle=p_bag_shuffle, ) mil2_head = mil2["mil_head"] be2 = mil2_head["bag_embed"] if be2 is None: raise RuntimeError("MIL level2 expected bag_embed (use attn + embed_cls).") level2_emb = be2 # [B, D] level2_logit_pool = self._logit_pool_from_mil_head(mil2_head) # [B, C] logits_level1 = self.head_level1(level1_emb) logits_level2 = self.head_level2(level2_emb) # --- Fusion: pool-logit gate -> keep fusion input dim unchanged --- if self.use_pool_gate_fusion: gate_in = torch.cat([level1_logit_pool, level2_logit_pool], dim=-1) # [B, 2C] g = torch.sigmoid(self.pool_gate(gate_in)) # [B, 1] fusion_in = torch.cat([g * level1_emb, (1.0 - g) * level2_emb], dim=-1) # [B, 2D] else: fusion_in = torch.cat([level1_emb, level2_emb], dim=-1) logits_fusion = self.fusion_head(fusion_in) out = { "logits_level1": logits_level1, "logits_level2": logits_level2, "logits_fusion": logits_fusion, } if return_viz: mil1_head = mil1.get("mil_head", {}) out["viz"] = { # Branch-level raw outputs "logits_level1": logits_level1, "logits_level2": logits_level2, "logits_fusion": logits_fusion, # MIL pooled logits and gate inputs "level1_logit_pool": level1_logit_pool, "level2_logit_pool": level2_logit_pool, # Embeddings before branch heads / fusion "level1_emb": level1_emb, "level2_emb": level2_emb, "fusion_in": fusion_in, # Attention maps and instance logits "tile_attn": tile_attn, "tile_inst_logits": tile_inst_logits, "level1_patch_attn": mil1_head.get("attn"), "level1_patch_inst_logits": mil1_head.get("inst_logits"), "level2_attn": mil2_head.get("attn"), "level2_inst_logits": mil2_head.get("inst_logits"), "num_tiles": torch.full( (b, 1), nt, device=logits_fusion.device, dtype=torch.int32 ), "num_patches_per_tile": torch.full( (b, 1), (t - 1), device=logits_fusion.device, dtype=torch.int32 ), # Scalar gate "gate": g if self.use_pool_gate_fusion else torch.full( (b, 1), 0.5, device=logits_fusion.device, dtype=logits_fusion.dtype ), "use_pool_gate_fusion": torch.tensor( [1 if self.use_pool_gate_fusion else 0], device=logits_fusion.device, dtype=logits_fusion.dtype, ), } return out class LTFNetV3_3(nn.Module): """Frozen CLIP visual + LTFCoreV3_3. Input: (B, N, 3, H, W).""" def __init__( self, clip_visual, num_heads=8, num_classes=2, num_layers=1, dropout=0.2, tile_heads=2, tile_layers=2, use_global_image=True, use_patch_transformer=True, linear_classifier=False, mil_heads=8, fusion_hidden=512, l1_cross_shuffle_p: float = 0.9, l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9), l1_with_replacement_p: float = 0.5, l2_cross_shuffle_p: float = 0.5, l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5), l2_with_replacement_p: float = 0.5, bag_token_shuffle_p: float = 0.5, cross_shuffle_min_keep_original: int = 0, use_mil_cross_shuffle: bool = True, use_clip_intermediate_for_mil: bool = False, clip_mil_stop_after_block_1based: int = 12, pool_gate_hidden: int = 16, use_pool_gate_fusion: bool = True, ): super().__init__() self.visual = clip_visual self.dtype = self.visual.conv1.weight.dtype for p in self.visual.parameters(): p.requires_grad = False self.visual.eval() self.input_resolution = self.visual.input_resolution self.embed_dim = self.visual.conv1.weight.shape[0] self.use_clip_intermediate_for_mil = bool(use_clip_intermediate_for_mil) self.clip_mil_stop_after_block_1based = int(clip_mil_stop_after_block_1based) if self.use_clip_intermediate_for_mil: n_bl = _count_clip_vit_depth(self.visual) stop = self.clip_mil_stop_after_block_1based if stop < 1 or stop > n_bl: raise ValueError( f"use_clip_intermediate_for_mil requires 1 <= clip_mil_stop_after_block_1based " f"<= {n_bl} (ViT depth), got {stop}" ) self.core = LTFCoreV3_3( embed_dim=self.embed_dim, num_heads=num_heads, num_classes=num_classes, num_layers=num_layers, dropout=dropout, tile_heads=tile_heads, tile_layers=tile_layers, use_global_image=use_global_image, use_patch_transformer=use_patch_transformer, linear_classifier=linear_classifier, mil_heads=mil_heads, fusion_hidden=fusion_hidden, l1_cross_shuffle_p=l1_cross_shuffle_p, l1_replace_ratio_range=l1_replace_ratio_range, l1_with_replacement_p=l1_with_replacement_p, l2_cross_shuffle_p=l2_cross_shuffle_p, l2_replace_ratio_range=l2_replace_ratio_range, l2_with_replacement_p=l2_with_replacement_p, bag_token_shuffle_p=bag_token_shuffle_p, cross_shuffle_min_keep_original=cross_shuffle_min_keep_original, use_mil_cross_shuffle=use_mil_cross_shuffle, pool_gate_hidden=pool_gate_hidden, use_pool_gate_fusion=use_pool_gate_fusion, ) def _stem_to_lnd(self, x_bchw: torch.Tensor) -> torch.Tensor: x = x_bchw.type(self.dtype) x = self.visual.conv1(x) x = x.reshape(x.shape[0], x.shape[1], -1) x = x.permute(0, 2, 1) cls = self.visual.class_embedding.to(x.dtype) + torch.zeros( x.shape[0], 1, x.shape[2], dtype=x.dtype, device=x.device ) x = torch.cat([cls, x], dim=1) x = x + self.visual.positional_embedding.to(x.dtype) x = self.visual.ln_pre(x) return x.permute(1, 0, 2) def _visual_tokens_from_imgs(self, imgs_bchw: torch.Tensor): x = self._stem_to_lnd(imgs_bchw) if not self.use_clip_intermediate_for_mil: x = self.visual.transformer(x) return x.permute(1, 0, 2), None blocks = _clip_vit_resblock_modules(self.visual.transformer) stop = self.clip_mil_stop_after_block_1based for i in range(stop): x = blocks[i](x) x_mil = x.permute(1, 0, 2).clone() for i in range(stop, len(blocks)): x = blocks[i](x) return x.permute(1, 0, 2), x_mil def forward( self, x: torch.Tensor, labels: Optional[torch.Tensor] = None, chunk_size=256, return_viz: bool = False, ): chunk_size = getattr(self, "eval_chunk_size", chunk_size) b, n, c, h, w = x.shape imgs = x.reshape(b * n, c, h, w) feats_chunks = [] mil_chunks = [] with torch.no_grad(): for start in range(0, imgs.size(0), chunk_size): end = min(start + chunk_size, imgs.size(0)) f_tok, m_tok = self._visual_tokens_from_imgs(imgs[start:end]) feats_chunks.append(f_tok) if m_tok is not None: mil_chunks.append(m_tok) feats_token = torch.cat(feats_chunks, dim=0) t = feats_token.shape[-2] d = feats_token.shape[-1] feats = feats_token.reshape(b, n, t, d).float() if not mil_chunks: return self.core(feats, labels=labels, x_mil=None, return_viz=return_viz) mil_token = torch.cat(mil_chunks, dim=0).reshape(b, n, t, d).float() return self.core(feats, labels=labels, x_mil=mil_token, return_viz=return_viz)