File size: 3,273 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 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 | """The classifier as a fused Linear with fixed ternary weights and one free bias.
model = FusedClassifier.from_hub()
score, present = model(image_tensor)
Constructed without a backbone the module still exposes `head`, the same
decision applied to an already-pooled vector.
"""
import argparse
import json
import sys
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
sys.path.insert(0, str(Path(__file__).resolve().parent)) # repo root, for `common`
from common import BACKBONE, pool # noqa: E402
HERE = Path(__file__).resolve().parent
class FusedClassifier(nn.Module):
"""Backbone -> 40-dim slice -> ternary linear head -> binary decision.
`retained_dims` indexes the 768-D pooled vector; `retained_weight` is +1 on
the person-positive positions and -1 on the person-negative ones. The
threshold is the only free parameter.
"""
def __init__(self, backbone, pos_dims, neg_dims, threshold):
super().__init__()
self.backbone = backbone
retained = list(pos_dims) + list(neg_dims)
self.register_buffer('retained_dims', torch.tensor(retained, dtype=torch.long))
w = torch.zeros(1, len(retained))
w[0, :len(pos_dims)] = 1.0
w[0, len(pos_dims):] = -1.0
self.register_buffer('retained_weight', w)
self.threshold = nn.Parameter(torch.tensor(float(threshold)))
def head(self, pooled):
"""(..., 768) pooled vector -> (score, present), no backbone involved."""
retained = pooled.index_select(-1, self.retained_dims)
score = F.linear(retained, self.retained_weight).squeeze(-1)
return score, score > self.threshold
@torch.inference_mode()
def forward(self, x):
"""x: (B, 3, 768, 768) normalized. Returns (score (B,), present (B,))."""
if self.backbone is None:
raise RuntimeError('constructed without a backbone; use .head(pooled)')
dev = 'cuda' if x.is_cuda else 'cpu'
with torch.autocast(dev, dtype=torch.bfloat16):
out = self.backbone.forward_features(x)
return self.head(pool(out['x_norm_patchtokens'].float()))
@classmethod
def from_config(cls, backbone=None, classifier_json=None):
c = json.loads(Path(classifier_json or HERE / 'classifier.json').read_text())
return cls(backbone, c['pos_dims'], c['neg_dims'], c['threshold'])
@classmethod
def from_hub(cls, repo_or_path=None, classifier_json=None):
from common.models import load_backbone
return cls.from_config(load_backbone(repo_or_path or BACKBONE), classifier_json)
if __name__ == '__main__':
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument('--classifier', type=Path, default=HERE / 'classifier.json')
args = ap.parse_args()
m = FusedClassifier.from_hub(classifier_json=args.classifier).eval()
n_all = sum(p.numel() for p in m.parameters())
n_backbone = sum(p.numel() for p in m.backbone.parameters())
print(f'total params: {n_all:,}')
print(f'backbone params: {n_backbone:,}')
print(f'head params: {n_all - n_backbone} '
f'(one learnable threshold; weights are fixed buffers)')
print(f'retained dims: {m.retained_dims.numel()}')
|