File size: 3,138 Bytes
bd659a9 | 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 | from typing import Dict, List, Optional, Tuple
import numpy as np
import torch
from sklearn.metrics import f1_score
from torch import nn
from torch.utils.data import DataLoader
from tqdm import tqdm
def move_batch(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
return {
"image": batch["image"].to(device=device, dtype=torch.float32),
"label": batch["label"].to(device=device, dtype=torch.float32),
"input_ids": batch["input_ids"].to(device=device, dtype=torch.long),
"attention_mask": batch["attention_mask"].to(device=device, dtype=torch.long),
"image_id": batch["image_id"].to(device=device),
}
def run_epoch(
model: nn.Module,
loader: DataLoader,
criterion: nn.Module,
device: torch.device,
optimizer: Optional[torch.optim.Optimizer] = None,
description: str = "Training",
) -> float:
is_training = optimizer is not None
model.train(is_training)
total_loss = 0.0
steps = 0
iterator = tqdm(loader, unit="batch", desc=description)
for batch in iterator:
batch = move_batch(batch, device)
if is_training:
optimizer.zero_grad(set_to_none=True)
logits = model(batch["input_ids"], batch["attention_mask"], batch["image"])
loss = criterion(logits, batch["label"])
if is_training:
loss.backward()
optimizer.step()
total_loss += float(loss.item())
steps += 1
iterator.set_postfix(loss=total_loss / max(steps, 1))
return total_loss / max(steps, 1)
@torch.no_grad()
def collect_probabilities(
model: nn.Module,
loader: DataLoader,
device: torch.device,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
model.eval()
probs: List[float] = []
labels: List[float] = []
image_ids: List[int] = []
for batch in loader:
batch = move_batch(batch, device)
logits = model(batch["input_ids"], batch["attention_mask"], batch["image"])
batch_probs = torch.sigmoid(logits)
probs.extend(batch_probs.detach().cpu().numpy().tolist())
labels.extend(batch["label"].detach().cpu().numpy().tolist())
image_ids.extend(batch["image_id"].detach().cpu().numpy().tolist())
return np.asarray(probs), np.asarray(labels), np.asarray(image_ids)
def threshold_grid(start: float, stop: float, step: float) -> List[float]:
values = []
current = start
while current <= stop + 1e-9:
values.append(round(current, 10))
current += step
return values
def optimize_threshold(
probs: np.ndarray,
labels: np.ndarray,
threshold_start: float,
threshold_stop: float,
threshold_step: float,
) -> Tuple[float, float, Dict[float, float]]:
scores: Dict[float, float] = {}
for threshold in threshold_grid(threshold_start, threshold_stop, threshold_step):
preds = (probs >= threshold).astype(float)
scores[threshold] = f1_score(labels, preds, average="macro", zero_division=0)
best_threshold = max(scores, key=scores.get)
return best_threshold, scores[best_threshold], scores
|