File size: 1,716 Bytes
e8b8483 | 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 | """Binary classification metrics, single-sourced so every stage scores identically."""
from typing import NamedTuple
import torch
class Metrics(NamedTuple):
"""F1, precision, recall, and the threshold they were measured at."""
f1: float
precision: float
recall: float
threshold: float = float('nan')
def asdict(self) -> dict:
d = {'F1': self.f1, 'precision': self.precision, 'recall': self.recall}
if self.threshold == self.threshold: # excludes NaN
d['threshold'] = self.threshold
return d
def prf1(pred: torch.Tensor, labels: torch.Tensor) -> Metrics:
"""Metrics for boolean prediction and label tensors."""
tp = (pred & labels).sum().float()
fp = (pred & ~labels).sum().float()
fn = (~pred & labels).sum().float()
precision = tp / (tp + fp).clamp(min=1)
recall = tp / (tp + fn).clamp(min=1)
f1 = 2 * precision * recall / (precision + recall).clamp(min=1e-9)
return Metrics(float(f1), float(precision), float(recall))
def f1_at(scores: torch.Tensor, labels: torch.Tensor, threshold: float) -> Metrics:
"""Metrics at a fixed threshold."""
return prf1(scores > threshold, labels)._replace(threshold=float(threshold))
def f1_sweep(scores: torch.Tensor, labels: torch.Tensor, n_candidates: int = 500) -> Metrics:
"""Best metrics over candidate thresholds drawn evenly from the sorted unique scores."""
uniq = torch.unique(scores).sort().values
stride = max(1, len(uniq) // n_candidates)
best = Metrics(0.0, 0.0, 0.0, 0.0)
for t in uniq.tolist()[::stride]:
m = prf1(scores > t, labels)
if m.f1 > best.f1:
best = m._replace(threshold=float(t))
return best
|