Spaces:
Sleeping
Sleeping
| 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 |