| """ |
| LTF V2: multi-branch (baseline + Level-1 MIL on patch tokens + Level-2 MIL on CLS + fusion). |
| V1 code in LTF.py is unchanged; this module is a separate copy for version V2. |
| """ |
| from typing import Optional, Tuple |
|
|
| import torch |
| import torch.nn as nn |
|
|
| from model.LTF import LTFCore |
| from mil_module.mil_core import MILCore |
|
|
|
|
| def _clip_vit_resblock_modules(transformer) -> list: |
| """Return the OpenAI CLIP ViT residual blocks as a list.""" |
| rb = transformer.resblocks |
| if isinstance(rb, nn.Sequential): |
| return list(rb.children()) |
| return list(rb) |
|
|
|
|
| def _count_clip_vit_depth(visual) -> int: |
| return len(_clip_vit_resblock_modules(visual.transformer)) |
|
|
|
|
| class BagTokenShuffle(nn.Module): |
| """Randomly permute tokens within each bag during training with probability p_shuffle.""" |
|
|
| def forward(self, x, p_shuffle: float = 0.0, p_drop: float = 0.0, key_padding_mask=None, **kwargs): |
| del p_drop, key_padding_mask, kwargs |
| if (not self.training) or p_shuffle <= 0: |
| return x |
| if float(torch.rand((), device=x.device)) >= float(p_shuffle): |
| return x |
| B, N, D = x.shape |
| if N <= 1: |
| return x |
| perm = torch.rand(B, N, device=x.device).argsort(dim=1) |
| return x.gather(1, perm.unsqueeze(-1).expand(B, N, D)) |
|
|
|
|
| class LTFCoreV2(nn.Module): |
| """Baseline and hierarchical MIL branches with shared frozen CLIP features. |
| |
| Level 1 pools patch tokens within local tiles and aggregates tile bags with a learned query. |
| Level 2 pools CLS tokens across views. Cross-Shuffle is applied only during training. |
| """ |
|
|
| 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, |
| ): |
| super().__init__() |
| self.embed_dim = embed_dim |
| self.use_global_image = use_global_image |
| self.num_classes = num_classes |
|
|
| 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.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, |
| ) |
|
|
| 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=True, |
| 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=True, |
| trans=BagTokenShuffle(), |
| allow_trans_token_drop=False, |
| ) |
|
|
| 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 |
| ) |
|
|
| self.head_level1 = nn.Linear(embed_dim, num_classes) |
| self.head_level2 = nn.Linear(embed_dim, num_classes) |
|
|
| classifier_in_dim = embed_dim * 2 if use_global_image else embed_dim |
| fusion_in = embed_dim + embed_dim + classifier_in_dim |
| self.fusion_head = nn.Sequential( |
| nn.Linear(fusion_in, fusion_hidden), |
| nn.ReLU(), |
| nn.Dropout(p=dropout), |
| nn.Linear(fusion_hidden, num_classes), |
| ) |
|
|
| 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): |
| """x: (B, N, T, D) -> same shape after optional token-level encoder.""" |
| B_2, N_2, T_2, D_2 = x.shape |
| xf = x.reshape(B_2 * N_2, T_2, D_2) |
| if self.baseline_core.patch_transformer is not None: |
| xf = self.baseline_core.patch_transformer(xf) |
| return xf.reshape(B_2, N_2, T_2, D_2) |
|
|
| def _forward_baseline_from_xfeat(self, x_feat): |
| """Run LTF classifier heads on processed tokens (same as LTFCore tail).""" |
| B_2, N_2, T_2, D_2 = x_feat.shape |
| bc = self.baseline_core |
| new_cls = x_feat[:, :, 0, :] |
| new_x = x_feat[:, :, 1:, :] |
| x_flat = new_x.reshape(-1, D_2) |
| x_result = bc.patch_classifier(x_flat) |
| x_result = x_result.reshape(B_2, N_2, T_2 - 1, 1) |
|
|
| main_feat = new_cls[:, 0] |
| tile_feats = new_cls[:, 1:] |
| cls_token = bc.cls_token.expand(B_2, 1, -1) |
| feature_tile = torch.cat([cls_token, tile_feats], dim=1) |
| feature_encoded = bc.tile_transformer(feature_tile) |
| cls_out = feature_encoded[:, 0] |
|
|
| if bc.use_global_image: |
| concat_feat = torch.cat([main_feat, cls_out], dim=1) |
| else: |
| concat_feat = cls_out |
| logits_b = bc.classifier(concat_feat) |
| jpeg_b = torch.sigmoid(bc.jpeg_estimate(cls_out)) |
| return logits_b, x_result, jpeg_b, concat_feat |
|
|
| def forward( |
| self, |
| x, |
| labels: Optional[torch.Tensor] = None, |
| x_mil: Optional[torch.Tensor] = None, |
| ): |
| """Run the branches on token features shaped [B, N, T, D]. |
| |
| N must include the global view and at least one local tile. Optional x_mil |
| provides alternate MIL tokens with the same shape. Binary labels enable |
| training-time Cross-Shuffle; the baseline uses the original x features. |
| """ |
| 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 |
| logits_baseline, patch_logits, jpeg_pred, concat_feat = self._forward_baseline_from_xfeat( |
| x_feat |
| ) |
|
|
| B, N, T, D = x_mil_feat.shape |
| if N < 2: |
| raise ValueError(f"LTFCoreV2 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 |
|
|
| 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) |
| q = self.level1_query.expand(B, -1, -1) |
| agg_out, _ = self.level1_aggregate_attn(q, bag_seq, bag_seq) |
| level1_emb = agg_out.squeeze(1) |
|
|
| cls_seq = x_mil_feat[:, :, 0, :] |
| 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, |
| ) |
| be2 = mil2["mil_head"]["bag_embed"] |
| if be2 is None: |
| raise RuntimeError("MIL level2 expected bag_embed (use attn + embed_cls).") |
| level2_emb = be2 |
|
|
| logits_level1 = self.head_level1(level1_emb) |
| logits_level2 = self.head_level2(level2_emb) |
|
|
| fusion_in = torch.cat([level1_emb, level2_emb, concat_feat], dim=-1) |
| logits_fusion = self.fusion_head(fusion_in) |
|
|
| return { |
| "logits_baseline": logits_baseline, |
| "logits_level1": logits_level1, |
| "logits_level2": logits_level2, |
| "logits_fusion": logits_fusion, |
| "patch_logits": patch_logits, |
| "jpeg_pred": jpeg_pred, |
| } |
|
|
|
|
| class LTFNetV2(nn.Module): |
| """Frozen CLIP visual + LTFCoreV2. 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_clip_intermediate_for_mil: bool = False, |
| clip_mil_stop_after_block_1based: int = 12, |
| ): |
| 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 <= v2_clip_mil_stop_after_block_1based " |
| f"<= {n_bl} (ViT depth), got {stop}" |
| ) |
|
|
| self.core = LTFCoreV2( |
| 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, |
| ) |
|
|
| def _stem_to_lnd(self, x_bchw: torch.Tensor) -> torch.Tensor: |
| """Convert images shaped [B, 3, H, W] to CLIP transformer input [L, N, D].""" |
| 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 |
| ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: |
| """Return final-layer tokens and optional intermediate-layer MIL tokens. |
| |
| The second result is None when intermediate MIL features are disabled. |
| """ |
| 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 get_tokens(self, x) -> torch.Tensor: |
| """Return final ViT tokens shaped [B, T, D], without separating intermediate features.""" |
| x = self._stem_to_lnd(x) |
| x = self.visual.transformer(x) |
| return x.permute(1, 0, 2) |
|
|
| def encode_image(self, image) -> torch.Tensor: |
| return self.visual(image.type(self.dtype)) |
|
|
| def forward(self, x: torch.Tensor, labels: Optional[torch.Tensor] = None, chunk_size=256): |
| 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) |
| mil_token = torch.cat(mil_chunks, dim=0).reshape(B, N, T, D).float() |
| return self.core(feats, labels=labels, x_mil=mil_token) |
|
|
