Translate legacy source documentation to English without changing executable code

0d8aa44 verified 1 day ago

7 kB

import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class MILHead(nn.Module):
    """Instance classification and MIL bag pooling.
    
    Supports max, log-sum-exp, top-k and gated-attention pooling. d_in is the
    token dimension; num_classes determines the classifier output size.
    Attention can pool instance logits or classify the pooled bag embedding.
    """
    def __init__(
        self,
        d_in: int = 1024,
        num_classes: int = 2,
        pooling: str = "lsep",
        lsep_temp: float = 1.0,
        lsep_norm: bool = True,
        topk: int = None,
        topk_ratio: float = 0.1,
        attn_hidden: int = 128,
        attn_dropout: float = 0.0,
    ):
        super().__init__()
        self.num_classes = num_classes
        self.pooling = pooling.lower()
        self.lsep_temp = float(lsep_temp)
        self.lsep_norm = bool(lsep_norm)
        self.topk = topk
        self.topk_ratio = float(topk_ratio)

        # per-token classifier (shared across tokens)
        self.classifier = nn.Linear(d_in, num_classes)

        # components for gated-attention MIL (Ilse et al., 2018)
        if self.pooling == "attn":
            print(f'Using gated-attention pooling with hidden size {attn_hidden}')
            self.attn_V = nn.Linear(d_in, attn_hidden, bias=True)
            self.attn_U = nn.Linear(d_in, attn_hidden, bias=True)
            self.attn_w = nn.Linear(attn_hidden, 1, bias=False)
            self.attn_drop = nn.Dropout(attn_dropout)

    @staticmethod
    def _masked_fill_min(x, mask, fill_value=None):
        """Utility: set masked positions to a very small number before max/topk."""
        if mask is None:
            return x
        if fill_value is None:
            fill_value = torch.finfo(x.dtype).min
        return x.masked_fill(~mask, fill_value)

    def _pool_max(self, logits, mask=None):
        # logits: [B, N, C] -> [B, C]
        logits = self._masked_fill_min(logits, mask.unsqueeze(-1) if mask is not None else None)
        bag = logits.max(dim=1).values
        return bag

    def _pool_lsep(self, logits, mask=None):
        # logits: [B, N, C] -> [B, C]
        # temperature T: divide logits by T, smaller T -> sharper
        z = logits / self.lsep_temp
        if mask is not None:
            
            neg_inf = torch.finfo(z.dtype).min
            z = z.masked_fill(~mask.unsqueeze(-1), neg_inf)
        bag = torch.logsumexp(z, dim=1)
        if self.lsep_norm:
            
            if mask is None:
                n = logits.size(1)
                bag = bag - math.log(max(n, 1))
            else:
                n_eff = mask.sum(dim=1).clamp_min(1).unsqueeze(-1)  # [B,1]
                bag = bag - n_eff.log()
        return bag

    def _pool_topk(self, logits, mask=None):
        # logits: [B, N, C] -> [B, C]
        B, N, C = logits.shape
        if mask is not None:
            logits = self._masked_fill_min(logits, mask.unsqueeze(-1))
            n_eff = mask.sum(dim=1)  # [B]
        else:
            n_eff = logits.new_full((B,), N)

        k = self.topk
        if k is None:
            # ratio -> at least 1
            k = torch.clamp((n_eff.float() * self.topk_ratio).long(), min=1)
        else:
            k = torch.clamp(torch.full_like(n_eff, k), min=1, max=N)

        
        
        
        scores = logits.transpose(1, 2)  # [B, C, N]
        
        topk_vals = []
        for b in range(B):
            
            kb = int(min(int(k[b].item()), int(n_eff[b].item())))
            vals, _ = scores[b].topk(kb, dim=-1)
            topk_vals.append(vals.mean(dim=-1, keepdim=True))  # [C,1]
        bag = torch.cat(topk_vals, dim=1).transpose(0,1).contiguous().squeeze(1)  # [B, C]
        return bag
    # v1
    # def _pool_attn(self, feats, logits, mask=None):
    #     """
    #     Gated-attention pooling:
    #         a_i ∝ w^T ( tanh(V h_i) ⊙ sigmoid(U h_i) )
    #         bag = Σ a_i * logits_i
    #     feats:  [B, N, D]
    #     logits: [B, N, C]
    #     mask:   [B, N] (True for valid)
    #     Returns:
    #         bag_logits: [B, C]
    #         attn:       [B, N, 1]
    #     """
    #     Vh = self.attn_V(feats)           # [B, N, L]
    #     Uh = self.attn_U(feats)           # [B, N, L]
    #     gate = torch.tanh(Vh) * torch.sigmoid(Uh)
    #     gate = self.attn_drop(gate)
    #     attn = self.attn_w(gate)          # [B, N, 1]

    #     if mask is not None:
    #         neg_inf = torch.finfo(attn.dtype).min
    #         attn = attn.masked_fill(~mask.unsqueeze(-1), neg_inf)

    #     attn = F.softmax(attn, dim=1)     # [B, N, 1]
    #     bag = (attn * logits).sum(dim=1)  # [B, C]
    #     return bag, attn
    
    def _pool_attn(self, feats, logits, mask=None):
        """Pool instance logits and features using normalized gated attention.
        
        The gate is tanh(Vh) times sigmoid(Uh), projected to one attention logit.
        Input features have shape [B, N, D]; instance logits are [B, N, C].
        The optional mask marks valid tokens. Return bag logits, attention and embedding.
        """
        Vh = self.attn_V(feats)           # [B, N, L]
        Uh = self.attn_U(feats)           # [B, N, L]
        gate = torch.tanh(Vh) * torch.sigmoid(Uh)
        gate = self.attn_drop(gate)
        attn = self.attn_w(gate)          # [B, N, 1]

        if mask is not None:
            neg_inf = torch.finfo(attn.dtype).min
            attn = attn.masked_fill(~mask.unsqueeze(-1), neg_inf)

        attn = F.softmax(attn, dim=1)        # [B, N, 1]

        
        bag_logits = (attn * logits).sum(dim=1)  # [B, C]

        
        bag_embed = (attn * feats).sum(dim=1)    # [B, D]

        return bag_logits, attn, bag_embed
    def forward(self, feats, mask=None):
        """Classify token features [B, N, D] and pool the bag.
        
        The optional [B, N] mask marks valid tokens. Return instance logits, bag
        logits, attention weights and a bag embedding when the selected mode provides them.
        """
        inst_logits = self.classifier(feats)  # [B, N, C]

        if self.pooling == "max":
            bag_logits = self._pool_max(inst_logits, mask)
            attn = None
        elif self.pooling == "lsep":
            bag_logits = self._pool_lsep(inst_logits, mask)
            attn = None
        elif self.pooling == "topk":
            bag_logits = self._pool_topk(inst_logits, mask)
            attn = None
        elif self.pooling == "attn":
            # v1
            # bag_logits, attn = self._pool_attn(feats, inst_logits, mask)
            # v2
            bag_logits, attn, bag_embed = self._pool_attn(feats, inst_logits, mask) 
        else:
            raise ValueError(f"Unknown pooling mode: {self.pooling}")

        return {
            "inst_logits": inst_logits,
            "bag_logits": bag_logits,
            "attn": attn,
            "bag_embed": bag_embed if self.pooling == "attn" else None,
        }

	import math
	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	class MILHead(nn.Module):
	"""Instance classification and MIL bag pooling.

	Supports max, log-sum-exp, top-k and gated-attention pooling. d_in is the
	token dimension; num_classes determines the classifier output size.
	Attention can pool instance logits or classify the pooled bag embedding.
	"""
	def __init__(
	self,
	d_in: int = 1024,
	num_classes: int = 2,
	pooling: str = "lsep",
	lsep_temp: float = 1.0,
	lsep_norm: bool = True,
	topk: int = None,
	topk_ratio: float = 0.1,
	attn_hidden: int = 128,
	attn_dropout: float = 0.0,
	):
	super().__init__()
	self.num_classes = num_classes
	self.pooling = pooling.lower()
	self.lsep_temp = float(lsep_temp)
	self.lsep_norm = bool(lsep_norm)
	self.topk = topk
	self.topk_ratio = float(topk_ratio)

	# per-token classifier (shared across tokens)
	self.classifier = nn.Linear(d_in, num_classes)

	# components for gated-attention MIL (Ilse et al., 2018)
	if self.pooling == "attn":
	print(f'Using gated-attention pooling with hidden size {attn_hidden}')
	self.attn_V = nn.Linear(d_in, attn_hidden, bias=True)
	self.attn_U = nn.Linear(d_in, attn_hidden, bias=True)
	self.attn_w = nn.Linear(attn_hidden, 1, bias=False)
	self.attn_drop = nn.Dropout(attn_dropout)

	@staticmethod
	def _masked_fill_min(x, mask, fill_value=None):
	"""Utility: set masked positions to a very small number before max/topk."""
	if mask is None:
	return x
	if fill_value is None:
	fill_value = torch.finfo(x.dtype).min
	return x.masked_fill(~mask, fill_value)

	def _pool_max(self, logits, mask=None):
	# logits: [B, N, C] -> [B, C]
	logits = self._masked_fill_min(logits, mask.unsqueeze(-1) if mask is not None else None)
	bag = logits.max(dim=1).values
	return bag

	def _pool_lsep(self, logits, mask=None):
	# logits: [B, N, C] -> [B, C]
	# temperature T: divide logits by T, smaller T -> sharper
	z = logits / self.lsep_temp
	if mask is not None:

	neg_inf = torch.finfo(z.dtype).min
	z = z.masked_fill(~mask.unsqueeze(-1), neg_inf)
	bag = torch.logsumexp(z, dim=1)
	if self.lsep_norm:

	if mask is None:
	n = logits.size(1)
	bag = bag - math.log(max(n, 1))
	else:
	n_eff = mask.sum(dim=1).clamp_min(1).unsqueeze(-1) # [B,1]
	bag = bag - n_eff.log()
	return bag

	def _pool_topk(self, logits, mask=None):
	# logits: [B, N, C] -> [B, C]
	B, N, C = logits.shape
	if mask is not None:
	logits = self._masked_fill_min(logits, mask.unsqueeze(-1))
	n_eff = mask.sum(dim=1) # [B]
	else:
	n_eff = logits.new_full((B,), N)

	k = self.topk
	if k is None:
	# ratio -> at least 1
	k = torch.clamp((n_eff.float() * self.topk_ratio).long(), min=1)
	else:
	k = torch.clamp(torch.full_like(n_eff, k), min=1, max=N)




	scores = logits.transpose(1, 2) # [B, C, N]

	topk_vals = []
	for b in range(B):

	kb = int(min(int(k[b].item()), int(n_eff[b].item())))
	vals, _ = scores[b].topk(kb, dim=-1)
	topk_vals.append(vals.mean(dim=-1, keepdim=True)) # [C,1]
	bag = torch.cat(topk_vals, dim=1).transpose(0,1).contiguous().squeeze(1) # [B, C]
	return bag
	# v1
	# def _pool_attn(self, feats, logits, mask=None):
	# """
	# Gated-attention pooling:
	# a_i ∝ w^T ( tanh(V h_i) ⊙ sigmoid(U h_i) )
	# bag = Σ a_i * logits_i
	# feats: [B, N, D]
	# logits: [B, N, C]
	# mask: [B, N] (True for valid)
	# Returns:
	# bag_logits: [B, C]
	# attn: [B, N, 1]
	# """
	# Vh = self.attn_V(feats) # [B, N, L]
	# Uh = self.attn_U(feats) # [B, N, L]
	# gate = torch.tanh(Vh) * torch.sigmoid(Uh)
	# gate = self.attn_drop(gate)
	# attn = self.attn_w(gate) # [B, N, 1]

	# if mask is not None:
	# neg_inf = torch.finfo(attn.dtype).min
	# attn = attn.masked_fill(~mask.unsqueeze(-1), neg_inf)

	# attn = F.softmax(attn, dim=1) # [B, N, 1]
	# bag = (attn * logits).sum(dim=1) # [B, C]
	# return bag, attn

	def _pool_attn(self, feats, logits, mask=None):
	"""Pool instance logits and features using normalized gated attention.

	The gate is tanh(Vh) times sigmoid(Uh), projected to one attention logit.
	Input features have shape [B, N, D]; instance logits are [B, N, C].
	The optional mask marks valid tokens. Return bag logits, attention and embedding.
	"""
	Vh = self.attn_V(feats) # [B, N, L]
	Uh = self.attn_U(feats) # [B, N, L]
	gate = torch.tanh(Vh) * torch.sigmoid(Uh)
	gate = self.attn_drop(gate)
	attn = self.attn_w(gate) # [B, N, 1]

	if mask is not None:
	neg_inf = torch.finfo(attn.dtype).min
	attn = attn.masked_fill(~mask.unsqueeze(-1), neg_inf)

	attn = F.softmax(attn, dim=1) # [B, N, 1]


	bag_logits = (attn * logits).sum(dim=1) # [B, C]


	bag_embed = (attn * feats).sum(dim=1) # [B, D]

	return bag_logits, attn, bag_embed
	def forward(self, feats, mask=None):
	"""Classify token features [B, N, D] and pool the bag.

	The optional [B, N] mask marks valid tokens. Return instance logits, bag
	logits, attention weights and a bag embedding when the selected mode provides them.
	"""
	inst_logits = self.classifier(feats) # [B, N, C]

	if self.pooling == "max":
	bag_logits = self._pool_max(inst_logits, mask)
	attn = None
	elif self.pooling == "lsep":
	bag_logits = self._pool_lsep(inst_logits, mask)
	attn = None
	elif self.pooling == "topk":
	bag_logits = self._pool_topk(inst_logits, mask)
	attn = None
	elif self.pooling == "attn":
	# v1
	# bag_logits, attn = self._pool_attn(feats, inst_logits, mask)
	# v2
	bag_logits, attn, bag_embed = self._pool_attn(feats, inst_logits, mask)
	else:
	raise ValueError(f"Unknown pooling mode: {self.pooling}")

	return {
	"inst_logits": inst_logits,
	"bag_logits": bag_logits,
	"attn": attn,
	"bag_embed": bag_embed if self.pooling == "attn" else None,
	}