Spaces:
Sleeping
Sleeping
File size: 3,271 Bytes
71c85a4 | 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 | import mlflow
import torch
from sklearn.metrics import average_precision_score, precision_score, recall_score
from src.models.evaluate import get_probs, metrics_at, best_threshold
def train_one_epoch(model, dataloader, loss_fn, optimizer, device):
model.train()
train_loss, correct, n_samples, n_batches = 0.0, 0, 0, 0
all_probs, all_y = [], []
for X, y in dataloader:
X, y = X.to(device), y.to(device)
logits = model(X)
loss = loss_fn(logits, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss += loss.item()
n_batches += 1
preds = (logits > 0).float()
correct += (preds == y).sum().item()
n_samples += y.numel()
all_probs.append(torch.sigmoid(logits).detach().cpu())
all_y.append(y.detach().cpu())
probs = torch.cat(all_probs).numpy().ravel()
y_true = torch.cat(all_y).numpy().ravel()
preds_np = (probs > 0.5).astype(int)
return {
"loss": train_loss / n_batches,
"acc": correct / n_samples,
"pr_auc": average_precision_score(y_true, probs),
"precision": precision_score(y_true, preds_np, zero_division=0),
"recall": recall_score(y_true, preds_np, zero_division=0),
}
def run_training(model, train_dl, test_dl, epochs, lr):
torch.manual_seed(42)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
# pos_weight = negatives / positives, from the TRAIN split
y_all = torch.cat([y for _, y in train_dl])
n_pos = y_all.sum().item()
n_neg = y_all.numel() - n_pos
pos_weight = torch.tensor([n_neg / max(n_pos, 1)], device=device)
print(f"pos_weight = {pos_weight.item():.1f} ({int(n_pos)} pos / {int(n_neg)} neg)")
loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=pos_weight)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
with mlflow.start_run():
for epoch in range(epochs):
m = train_one_epoch(model, train_dl, loss_fn, optimizer, device)
for k, v in m.items():
mlflow.log_metric(f"train_{k}", v, step=epoch)
print(f"epoch {epoch}: loss={m['loss']:.4f} acc={m['acc']:.4f} "
f"pr_auc={m['pr_auc']:.4f} P={m['precision']:.3f} R={m['recall']:.3f}")
# ---------- evaluate on the held-out TEST set ----------
probs, y_true = get_probs(model, test_dl)
# (a) the naive 0.5 cutoff
base = metrics_at(probs, y_true, 0.5)
print(f"\nTEST @0.50 : P={base['precision']:.3f} R={base['recall']:.3f} "
f"F1={base['f1']:.3f} PR-AUC={base['pr_auc']:.3f}")
# (b) the F1-best cutoff — same model, better trade-off
thr = best_threshold(probs, y_true)
tuned = metrics_at(probs, y_true, thr)
print(f"TEST @{thr:.2f} : P={tuned['precision']:.3f} R={tuned['recall']:.3f} "
f"F1={tuned['f1']:.3f}")
# log BOTH so you can compare in MLflow
for k, v in base.items():
mlflow.log_metric(f"test_{k}", v)
for k, v in tuned.items():
mlflow.log_metric(f"test_tuned_{k}", v)
return model |