Release original LTF v3.3 epoch 20 inference model

acb2ce3 verified 2 days ago

18.5 kB

"""
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)

	"""
	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 = [glevel1_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)