File size: 2,064 Bytes
f5498f9 | 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 | """Pooled feature extraction and the ternary scoring rule.
The classifier reads one 768-D vector per image: layernorm across the 768
channels of every patch token, then max-pool across the 2304 patches. Score is
the sum of the person-positive dims minus the sum of the person-negative dims.
"""
from typing import Optional, Sequence
import torch
import torch.nn.functional as F
D = 768
RES = 768
def pool(patch_tokens: torch.Tensor) -> torch.Tensor:
"""(N, D) or (B, N, D) patch tokens -> (D,) or (B, D) pooled vector."""
ln = F.layer_norm(patch_tokens.float(), [D])
return ln.max(dim=-2).values
@torch.inference_mode()
def backbone_pooled(backbone, x: torch.Tensor, autocast: bool = True) -> torch.Tensor:
"""Forward a normalized batch through the backbone and pool it."""
if autocast:
dev = 'cuda' if x.is_cuda else 'cpu'
with torch.autocast(dev, dtype=torch.bfloat16):
out = backbone.forward_features(x)
else:
out = backbone.forward_features(x)
return pool(out['x_norm_patchtokens'].float())
def score(pooled: torch.Tensor, pos: Sequence[int], neg: Sequence[int]) -> torch.Tensor:
"""sum(pooled[pos]) - sum(pooled[neg]), over the last axis."""
if not torch.is_tensor(pos):
pos = torch.tensor(list(pos), dtype=torch.long, device=pooled.device)
if not torch.is_tensor(neg):
neg = torch.tensor(list(neg), dtype=torch.long, device=pooled.device)
return pooled.index_select(-1, pos).sum(-1) - pooled.index_select(-1, neg).sum(-1)
def score_pool(backbone, loaded, pos, neg, target_dims: Optional[torch.Tensor] = None):
"""Score a pool; with `target_dims`, also return the pooled activations there."""
scores, targets = [], []
for x in loaded:
pooled = backbone_pooled(backbone, x)[0]
scores.append(score(pooled, pos, neg))
if target_dims is not None:
targets.append(pooled[target_dims])
stacked = torch.stack(scores)
return (stacked, torch.stack(targets)) if target_dims is not None else (stacked, None)
|