source/train.py

from __future__ import annotations

import json
import random
from pathlib import Path

import numpy as np
import torch
import trackio
from model import MeshGraphGCN, normalize_adjacency, parameter_count
from safetensors.torch import save_file
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    average_precision_score,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
    roc_auc_score,
)
from torch.nn import functional as F

PROJECT_DIR = Path(__file__).resolve().parent
DATA_DIR = PROJECT_DIR / "data"
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "meshgraph-gcn"


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def metrics(labels: np.ndarray, scores: np.ndarray, threshold: float = 0.5) -> dict:
    predictions = scores >= threshold
    return {
        "accuracy": float(accuracy_score(labels, predictions)),
        "roc_auc": float(roc_auc_score(labels, scores)),
        "average_precision": float(average_precision_score(labels, scores)),
        "precision": float(precision_score(labels, predictions, zero_division=0)),
        "recall": float(recall_score(labels, predictions, zero_division=0)),
        "f1": float(f1_score(labels, predictions, zero_division=0)),
        "confusion_matrix": confusion_matrix(labels, predictions).tolist(),
    }


def best_threshold(labels: np.ndarray, scores: np.ndarray) -> float:
    candidates = np.quantile(scores, np.linspace(0.05, 0.95, 300))
    return float(
        max(
            candidates,
            key=lambda threshold: f1_score(
                labels,
                scores >= threshold,
                zero_division=0,
            ),
        )
    )


def main() -> None:
    seed_everything(2033)
    graph = np.load(DATA_DIR / "meshgraph.npz")
    features = graph["features"].astype(np.float32)
    labels = graph["labels"].astype(np.int64)
    train_indices = graph["train_indices"]
    validation_indices = graph["validation_indices"]
    test_indices = graph["test_indices"]
    mean = features[train_indices].mean(axis=0)
    scale = np.maximum(features[train_indices].std(axis=0), 1e-5)
    features = (features - mean) / scale
    feature_tensor = torch.from_numpy(features)
    label_tensor = torch.from_numpy(labels)
    adjacency = torch.from_numpy(graph["adjacency"].astype(np.float32))
    normalized = normalize_adjacency(adjacency)

    baseline = LogisticRegression(
        class_weight="balanced",
        max_iter=2000,
        random_state=2033,
    )
    baseline.fit(features[train_indices], labels[train_indices])
    baseline_validation = baseline.predict_proba(features[validation_indices])[:, 1]
    baseline_threshold = best_threshold(
        labels[validation_indices],
        baseline_validation,
    )
    baseline_test = baseline.predict_proba(features[test_indices])[:, 1]

    model = MeshGraphGCN(features=features.shape[1])
    positive_weight = (labels[train_indices] == 0).sum() / max(
        1, (labels[train_indices] == 1).sum()
    )
    class_weights = torch.tensor([1.0, positive_weight], dtype=torch.float32)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.01, weight_decay=0.002)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=350)
    best_validation_auc = -1.0
    best_epoch = 0
    best_state = None
    trackio.init(
        project="meshgraph-gcn",
        name="two-layer-gcn-v1",
        config={
            "parameters": parameter_count(model),
            "nodes": len(labels),
            "edges": int(adjacency.sum().item() // 2),
            "train_labels": len(train_indices),
            "transductive": True,
        },
    )
    for epoch in range(1, 351):
        model.train()
        logits = model(feature_tensor, normalized)
        loss = F.cross_entropy(
            logits[train_indices],
            label_tensor[train_indices],
            weight=class_weights,
        )
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
        scheduler.step()
        model.eval()
        with torch.no_grad():
            probabilities = model(feature_tensor, normalized).softmax(dim=1)[:, 1]
        validation_auc = roc_auc_score(
            labels[validation_indices],
            probabilities[validation_indices].numpy(),
        )
        if validation_auc > best_validation_auc:
            best_validation_auc = validation_auc
            best_epoch = epoch
            best_state = {
                key: value.detach().cpu().clone()
                for key, value in model.state_dict().items()
            }
        if epoch == 1 or epoch % 10 == 0:
            trackio.log(
                {
                    "epoch": epoch,
                    "train_loss": float(loss.detach()),
                    "validation_roc_auc": validation_auc,
                    "learning_rate": scheduler.get_last_lr()[0],
                }
            )
    trackio.finish()
    assert best_state is not None
    model.load_state_dict(best_state)
    model.eval()
    with torch.no_grad():
        gcn_scores = model(feature_tensor, normalized).softmax(dim=1)[:, 1].numpy()
    gcn_threshold = best_threshold(
        labels[validation_indices],
        gcn_scores[validation_indices],
    )
    results = {
        "model": "MeshGraph GCN",
        "parameters": parameter_count(model),
        "nodes": len(labels),
        "edges": int(adjacency.sum().item() // 2),
        "best_epoch": best_epoch,
        "best_validation_roc_auc": best_validation_auc,
        "gcn_threshold": gcn_threshold,
        "gcn_test": metrics(
            labels[test_indices],
            gcn_scores[test_indices],
            gcn_threshold,
        ),
        "feature_only_logistic_test": metrics(
            labels[test_indices],
            baseline_test,
            baseline_threshold,
        ),
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    save_file(model.state_dict(), ARTIFACT_DIR / "model.safetensors")
    np.savez(
        ARTIFACT_DIR / "preprocessing.npz",
        mean=mean,
        scale=scale,
    )
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()