File size: 6,997 Bytes
acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 0d8aa44 acb2ce3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 | 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,
} |