run.py

"""End-to-end runner for Project 2.

Running `python run.py` does the following in one shot:

  Tasks 1-4 (source organ):
    - Grid search over learning rate, batch size, and dropout.
    - Each combination trains a fresh model and logs Loss/train, Loss/test,
      Accuracy/test to TensorBoard every epoch.
    - The model graph is logged with writer.add_graph (Task 1).
    - Final accuracy/loss/precision/recall are written to the HParams tab (Task 2).
    - A confusion matrix figure is added to TensorBoard (Task 3).
    - The best-by-test-accuracy state_dict from each run is saved, and the overall
      best run's checkpoint is promoted to results/best_dna_cnn.pth (Task 4).

  Task 5 (target organ):
    - Loads results/best_dna_cnn.pth, freezes the convolutional backbone,
      reinitializes the classifier head, and fine-tunes for 5 epochs on a
      different organ. Everything is logged to TensorBoard.

All knobs live in the CONFIG dict at the top of this file.
"""

import copy
import itertools
import json
from pathlib import Path

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from sklearn.metrics import (
    ConfusionMatrixDisplay,
    confusion_matrix,
    precision_score,
    recall_score,
)
from torch.optim import Adam
from torch.utils.tensorboard import SummaryWriter

from dataset import create_loaders, infer_seq_len_from_file
from model import DNAAccessibilityCNN


CONFIG = {
    # Source organ data (used for tasks 1-4).
    "source_positive": "data/all_positive.txt",
    "source_negative": "data/all_negative.txt",

    # Target organ data (used for task 5 - frozen-backbone transfer learning).
    "target_positive": "data/brain/positive.txt",
    "target_negative": "data/brain/negative.txt",

    # Training schedule.
    "epochs": 10,
    "transfer_epochs": 5,

    # Hyperparameter grid (task 2 - 3 hyperparameters).
    "learning_rates": [1e-3, 1e-4],
    "batch_sizes": [32, 64],
    "dropouts": [0.3, 0.5],

    # Fixed model hyperparameters.
    "kernel_size": 5,
    "hidden_dim": 256,

    # Output locations.
    "runs_dir": "runs",
    "results_dir": "results",
    "best_checkpoint_name": "best_dna_cnn.pth",

    # Reproducibility.
    "seed": 42,
}


def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    running_loss = 0.0
    seen = 0
    for x, y in loader:
        x = x.to(device)
        y = y.to(device)

        optimizer.zero_grad()
        probs = model(x)
        loss = criterion(probs, y)
        loss.backward()
        optimizer.step()

        running_loss += loss.item() * y.size(0)
        seen += y.size(0)

    return running_loss / seen


def evaluate(model, loader, criterion, device):
    """Return (avg_loss, accuracy, all_labels, all_preds) on the given loader."""
    model.eval()
    running_loss = 0.0
    correct = 0
    seen = 0
    all_labels = []
    all_preds = []

    with torch.no_grad():
        for x, y in loader:
            x = x.to(device)
            y = y.to(device)

            probs = model(x)
            loss = criterion(probs, y)
            preds = (probs >= 0.5).float()

            running_loss += loss.item() * y.size(0)
            correct += (preds == y).sum().item()
            seen += y.size(0)

            all_labels.extend(y.int().cpu().tolist())
            all_preds.extend(preds.int().cpu().tolist())

    return running_loss / seen, correct / seen, all_labels, all_preds


def make_confusion_figure(labels, preds, title):
    cm = confusion_matrix(labels, preds, labels=[0, 1])
    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=["Negative", "Positive"])
    fig, ax = plt.subplots()
    disp.plot(ax=ax, values_format="d")
    ax.set_title(title)
    fig.tight_layout()
    return fig


def train_full_run(model, train_loader, test_loader, epochs, learning_rate, device, writer, seq_len, run_checkpoint):
    """Train for `epochs` epochs, log to TensorBoard, save best-by-test-accuracy checkpoint.

    Returns a dict of final metrics (best test accuracy, final test loss, precision, recall).
    """
    criterion = nn.BCELoss()
    optimizer = Adam(model.parameters(), lr=learning_rate)

    # Task 1: log the model graph.
    dummy_input = torch.zeros(1, 4, seq_len, device=device)
    writer.add_graph(model, dummy_input)

    best_test_acc = -1.0
    best_state = None

    for epoch in range(1, epochs + 1):
        train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device)
        test_loss, test_acc, _, _ = evaluate(model, test_loader, criterion, device)

        # Task 1: per-epoch scalar logs.
        writer.add_scalar("Loss/train", train_loss, epoch)
        writer.add_scalar("Loss/test", test_loss, epoch)
        writer.add_scalar("Accuracy/test", test_acc, epoch)

        print(f"  epoch {epoch:02d}/{epochs} | train_loss={train_loss:.4f} | test_loss={test_loss:.4f} | test_acc={test_acc:.4f}")

        # Task 4: keep the best-by-test-accuracy state in memory and on disk.
        if test_acc > best_test_acc:
            best_test_acc = test_acc
            best_state = copy.deepcopy(model.state_dict())
            torch.save(best_state, run_checkpoint)

    # Reload the best checkpoint so all final reporting matches the saved model.
    model.load_state_dict(best_state)

    test_loss, test_acc, labels, preds = evaluate(model, test_loader, criterion, device)
    precision = precision_score(labels, preds, zero_division=0)
    recall = recall_score(labels, preds, zero_division=0)

    # Task 3: log confusion matrix figure plus precision and recall.
    fig = make_confusion_figure(labels, preds, "Confusion Matrix")
    writer.add_figure("ConfusionMatrix", fig, global_step=epochs)
    writer.add_scalar("Metrics/precision", precision, epochs)
    writer.add_scalar("Metrics/recall", recall, epochs)
    plt.close(fig)

    return {
        "best_test_acc": best_test_acc,
        "test_loss": test_loss,
        "test_accuracy": test_acc,
        "precision": precision,
        "recall": recall,
    }


def run_grid_search(config, seq_len, device):
    """Tasks 1-4. Sweep the (lr, batch_size, dropout) grid and promote the global best checkpoint."""
    runs_dir = Path(config["runs_dir"])
    results_dir = Path(config["results_dir"])
    runs_dir.mkdir(parents=True, exist_ok=True)
    results_dir.mkdir(parents=True, exist_ok=True)

    best_global_acc = -1.0
    best_config = None
    best_run_checkpoint = None
    summary = []

    grid = list(itertools.product(
        config["learning_rates"],
        config["batch_sizes"],
        config["dropouts"],
    ))
    print(f"\nGrid search: {len(grid)} combinations, {config['epochs']} epochs each.\n")

    for lr, batch_size, dropout in grid:
        torch.manual_seed(config["seed"])

        run_name = f"lr_{lr}_bs_{batch_size}_drop_{dropout}"
        run_dir = runs_dir / run_name
        run_dir.mkdir(parents=True, exist_ok=True)
        run_checkpoint = run_dir / "best_state_dict.pth"

        print(f"=== {run_name} ===")

        train_loader, test_loader = create_loaders(
            positive_file=config["source_positive"],
            negative_file=config["source_negative"],
            batch_size=batch_size,
            seq_len=seq_len,
            seed=config["seed"],
        )

        model = DNAAccessibilityCNN(
            seq_len=seq_len,
            kernel_size=config["kernel_size"],
            hidden_dim=config["hidden_dim"],
            dropout=dropout,
        ).to(device)

        writer = SummaryWriter(log_dir=str(run_dir))
        metrics = train_full_run(
            model=model,
            train_loader=train_loader,
            test_loader=test_loader,
            epochs=config["epochs"],
            learning_rate=lr,
            device=device,
            writer=writer,
            seq_len=seq_len,
            run_checkpoint=run_checkpoint,
        )

        # Task 2: log final metrics to the HParams tab so all combos can be compared side-by-side.
        writer.add_hparams(
            {"lr": lr, "batch_size": batch_size, "dropout": dropout},
            {
                "hparam/test_acc": metrics["test_accuracy"],
                "hparam/test_loss": metrics["test_loss"],
                "hparam/precision": metrics["precision"],
                "hparam/recall": metrics["recall"],
            },
        )
        writer.close()

        record = {
            "run_name": run_name,
            "learning_rate": lr,
            "batch_size": batch_size,
            "dropout": dropout,
        }
        record.update(metrics)
        summary.append(record)

        if metrics["best_test_acc"] > best_global_acc:
            best_global_acc = metrics["best_test_acc"]
            best_config = {"learning_rate": lr, "batch_size": batch_size, "dropout": dropout}
            best_run_checkpoint = run_checkpoint

    # Task 4: promote the overall best run's state_dict to the canonical path.
    best_checkpoint_path = results_dir / config["best_checkpoint_name"]
    best_state = torch.load(best_run_checkpoint, map_location="cpu")
    torch.save(best_state, best_checkpoint_path)

    with open(results_dir / "grid_summary.json", "w") as f:
        json.dump(summary, f, indent=2)

    print(f"\nBest grid run: {best_config} -> test_acc={best_global_acc:.4f}")
    print(f"Saved best checkpoint to: {best_checkpoint_path}")
    return best_config, best_checkpoint_path


def run_transfer_learning(config, seq_len, device, source_config, checkpoint_path):
    """Task 5. Frozen-backbone transfer learning on a different organ."""
    runs_dir = Path(config["runs_dir"])
    results_dir = Path(config["results_dir"])

    train_loader, test_loader = create_loaders(
        positive_file=config["target_positive"],
        negative_file=config["target_negative"],
        batch_size=source_config["batch_size"],
        seq_len=seq_len,
        seed=config["seed"],
    )

    # Step 2 (PDF): rebuild the model and load the saved weights.
    model = DNAAccessibilityCNN(
        seq_len=seq_len,
        kernel_size=config["kernel_size"],
        hidden_dim=config["hidden_dim"],
        dropout=source_config["dropout"],
    )
    state = torch.load(checkpoint_path, map_location="cpu")
    model.load_state_dict(state)

    total_before = 0
    for p in model.parameters():
        if p.requires_grad:
            total_before += p.numel()
    print(f"\nTrainable params before freezing: {total_before}")

    # Step 3: freeze the convolutional backbone.
    for p in model.features.parameters():
        p.requires_grad = False

    # Step 4: replace the classifier so the head starts fresh on the new organ.
    model.classifier = nn.Sequential(
        nn.Flatten(),
        nn.Linear(model.feature_dim, config["hidden_dim"]),
        nn.ReLU(),
        nn.Dropout(source_config["dropout"]),
        nn.Linear(config["hidden_dim"], 1),
        nn.Sigmoid(),
    )
    model = model.to(device)

    trainable_params = []
    for p in model.parameters():
        if p.requires_grad:
            trainable_params.append(p)
    total_after = sum(p.numel() for p in trainable_params)
    print(f"Trainable params after freezing (head only): {total_after}")

    criterion = nn.BCELoss()
    optimizer = Adam(trainable_params, lr=source_config["learning_rate"])

    writer = SummaryWriter(log_dir=str(runs_dir / "transfer"))
    dummy_input = torch.zeros(1, 4, seq_len, device=device)
    writer.add_graph(model, dummy_input)

    epochs = config["transfer_epochs"]
    print(f"Fine-tuning for {epochs} epochs (head only).")
    for epoch in range(1, epochs + 1):
        train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device)
        test_loss, test_acc, _, _ = evaluate(model, test_loader, criterion, device)
        writer.add_scalar("Loss/train", train_loss, epoch)
        writer.add_scalar("Loss/test", test_loss, epoch)
        writer.add_scalar("Accuracy/test", test_acc, epoch)
        print(f"  transfer epoch {epoch:02d}/{epochs} | train_loss={train_loss:.4f} | test_loss={test_loss:.4f} | test_acc={test_acc:.4f}")

    # Final evaluation, confusion matrix, and HParams entry for the transfer run.
    test_loss, test_acc, labels, preds = evaluate(model, test_loader, criterion, device)
    precision = precision_score(labels, preds, zero_division=0)
    recall = recall_score(labels, preds, zero_division=0)

    fig = make_confusion_figure(labels, preds, "Transfer Confusion Matrix")
    writer.add_figure("ConfusionMatrix", fig, global_step=epochs)
    writer.add_scalar("Metrics/precision", precision, epochs)
    writer.add_scalar("Metrics/recall", recall, epochs)
    writer.add_hparams(
        {
            "lr": source_config["learning_rate"],
            "batch_size": source_config["batch_size"],
            "dropout": source_config["dropout"],
            "frozen_backbone": 1,
        },
        {
            "hparam/test_acc": test_acc,
            "hparam/test_loss": test_loss,
            "hparam/precision": precision,
            "hparam/recall": recall,
        },
    )
    writer.close()
    plt.close(fig)

    transfer_metrics = {
        "test_loss": test_loss,
        "test_accuracy": test_acc,
        "precision": precision,
        "recall": recall,
        "source_config": source_config,
    }
    with open(results_dir / "transfer_metrics.json", "w") as f:
        json.dump(transfer_metrics, f, indent=2)

    torch.save(model.state_dict(), results_dir / "transfer_model.pth")
    print(f"Transfer test accuracy: {test_acc:.4f} | precision: {precision:.4f} | recall: {recall:.4f}")


def main():
    torch.manual_seed(CONFIG["seed"])

    seq_len = infer_seq_len_from_file(CONFIG["source_positive"])
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device} | seq_len: {seq_len}")

    print("\n--- Tasks 1-4: grid search on the source organ ---")
    best_config, best_checkpoint_path = run_grid_search(CONFIG, seq_len, device)

    print("\n--- Task 5: frozen-backbone transfer learning on the target organ ---")
    run_transfer_learning(CONFIG, seq_len, device, best_config, best_checkpoint_path)


if __name__ == "__main__":
    main()