alnn_rewrite/model/training.py

import pathlib as pl
import time
from threading import Event
from typing import Callable, Tuple, cast

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from model.dataset import ADNIDataset
from util.progress import ProgressTracker
from util.ui_logger import PipelineLogger

type TrainMetrics = Tuple[
    float, float, float, float
]  # (train_loss, val_loss, train_acc, val_acc)

type TestMetrics = Tuple[float, float]  # (test_loss, test_acc)
type KLLossFn = Callable[[nn.Module], torch.Tensor | None]


def _model_device(model: nn.Module) -> torch.device:
    first_param = next(model.parameters(), None)
    if first_param is None:
        return torch.device("cpu")
    return first_param.device


def _move_batch_to_model_device(
    model: nn.Module,
    mri: torch.Tensor,
    xls: torch.Tensor,
    targets: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    device = _model_device(model)
    mri = mri.to(device, non_blocking=True)
    xls = xls.to(device, non_blocking=True)
    targets = targets.to(device, non_blocking=True)
    return mri, xls, targets


def _batch_correct_and_total(
    outputs: torch.Tensor,
    targets: torch.Tensor,
) -> Tuple[int, int]:
    if outputs.ndim > 1 and outputs.size(-1) > 1:
        predicted_classes = outputs.argmax(dim=1)
        if targets.ndim > 1 and targets.size(-1) > 1:
            target_classes = targets.argmax(dim=1)
        else:
            target_classes = targets.long().view(-1)

        correct = int((predicted_classes == target_classes).sum().item())
        total = int(target_classes.numel())
        return correct, total

    predicted = (outputs > 0.5).to(dtype=targets.dtype)
    correct = int((predicted == targets).sum().item())
    total = int(targets.numel())
    return correct, total


def _check_control_events(
    stop_event: Event | None,
    pause_event: Event | None,
) -> None:
    if stop_event is not None and stop_event.is_set():
        raise InterruptedError("Pipeline execution stopped by user.")
    while pause_event is not None and pause_event.is_set():
        time.sleep(0.5)
        if stop_event is not None and stop_event.is_set():
            raise InterruptedError("Pipeline execution stopped by user while paused.")


def test_model(
    model: nn.Module,
    test_loader: DataLoader[ADNIDataset],
    criterion: nn.Module,
    progress: ProgressTracker,
    log: PipelineLogger,
    stop_event: Event | None = None,
    pause_event: Event | None = None,
) -> TestMetrics:
    """
    Tests the model on the test dataset.

    Args:
        model (nn.Module): The model to test.
        test_loader (DataLoader[ADNIDataset]): DataLoader for the test dataset.
        criterion (nn.Module): Loss function to compute the loss.

    Returns:
        TrainMetrics: A tuple containing the test loss and test accuracy.
    """
    model.eval()
    test_loss = 0.0
    correct = 0
    total = 0
    total_samples = 0

    test_progress = progress.get_sub_tracker("Testing Batches")
    test_progress.update(total=len(test_loader), advance=0)
    with torch.no_grad():
        for _, (mri, xls, targets, _) in enumerate(test_loader):
            _check_control_events(stop_event=stop_event, pause_event=pause_event)
            mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
            outputs = model((mri, xls))
            loss = criterion(outputs, targets)
            batch_size = mri.size(0)
            test_loss += loss.item() * batch_size
            total_samples += batch_size

            # Calculate accuracy
            batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
            correct += batch_correct
            total += batch_total
            test_progress.update(total=len(test_loader), advance=1)

    test_loss = test_loss / total_samples if total_samples > 0 else 0.0
    test_acc = correct / total if total > 0 else 0.0
    return test_loss, test_acc


def train_epoch(
    model: nn.Module,
    train_loader: DataLoader[ADNIDataset],
    val_loader: DataLoader[ADNIDataset],
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    log: PipelineLogger,
    progress: ProgressTracker,
    stop_event: Event | None = None,
    pause_event: Event | None = None,
) -> Tuple[float, float, float, float]:
    """
    Trains the model for one epoch and evaluates it on the validation set.

    Args:
        model (nn.Module): The model to train.
        train_loader (DataLoader[ADNIDataset]): DataLoader for the training dataset.
        val_loader (DataLoader[ADNIDataset]): DataLoader for the validation dataset.
        optimizer (torch.optim.Optimizer): Optimizer for updating model parameters.
        criterion (nn.Module): Loss function to compute the loss.

    Returns:
        Tuple[float, float, float, float]: A tuple containing the training loss, validation loss, training accuracy, and validation accuracy.
    """
    model.train()
    train_loss = 0.0
    train_correct = 0
    train_total = 0
    train_total_samples = 0

    # Training loop - with progress bar for the training batches
    train_progress = progress.get_sub_tracker("Training Batches")
    train_progress.reset()
    train_progress.set_title("Training Batches")
    train_progress.update(total=len(train_loader), advance=0)
    for _, (mri, xls, targets, _) in enumerate(train_loader):
        _check_control_events(stop_event=stop_event, pause_event=pause_event)
        mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
        optimizer.zero_grad()
        outputs = model((mri, xls))
        loss = criterion(outputs, targets)
        loss.backward()  # type: ignore[reportUnknownMemberType]
        optimizer.step()
        batch_size = mri.size(0)
        train_loss += loss.item() * batch_size
        train_total_samples += batch_size
        batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
        train_correct += batch_correct
        train_total += batch_total
        train_progress.update(total=len(train_loader), advance=1)
    train_loss = train_loss / train_total_samples if train_total_samples > 0 else 0.0

    model.eval()
    val_loss = 0.0
    correct = 0
    total = 0
    val_total_samples = 0

    test_progress = progress.get_sub_tracker("Validation Batches")
    test_progress.reset()
    test_progress.set_title("Validation Batches")
    test_progress.update(total=len(val_loader), advance=0)
    with torch.no_grad():
        for _, (mri, xls, targets, _) in enumerate(val_loader):
            _check_control_events(stop_event=stop_event, pause_event=pause_event)
            mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
            outputs = model((mri, xls))
            loss = criterion(outputs, targets)
            batch_size = mri.size(0)
            val_loss += loss.item() * batch_size
            val_total_samples += batch_size

            # Calculate accuracy
            batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
            correct += batch_correct
            total += batch_total
            test_progress.update(total=len(val_loader), advance=1)

    val_loss = val_loss / val_total_samples if val_total_samples > 0 else 0.0
    val_acc = correct / total if total > 0 else 0.0
    train_acc = train_correct / train_total if train_total > 0 else 0.0
    return train_loss, val_loss, train_acc, val_acc


def train_model(
    model: nn.Module,
    train_loader: DataLoader[ADNIDataset],
    val_loader: DataLoader[ADNIDataset],
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    num_epochs: int,
    output_path: pl.Path,
    progress: ProgressTracker,
    log: PipelineLogger,
    stop_event: Event | None = None,
    pause_event: Event | None = None,
) -> Tuple[nn.Module, pd.DataFrame]:
    """
    Trains the model using the provided training and validation data loaders.

    Args:
        model (nn.Module): The model to train.
        train_loader (DataLoader[ADNIDataset]): DataLoader for the training dataset.
        val_loader (DataLoader[ADNIDataset]): DataLoader for the validation dataset.
        num_epochs (int): Number of epochs to train the model.
        learning_rate (float): Learning rate for the optimizer.

    Returns:
        Result[nn.Module, str]: A Result object containing the trained model or an error message.
    """

    # Record the training history
    # We record the Epoch, Training Loss, Validation Loss, Training Accuracy, and Validation Accuracy
    # use a (num_epochs, 4) shape ndarray to store the history before creating the DataArray

    nhist = np.zeros((num_epochs, 4), dtype=np.float32)

    # Get new progress bar for epochs

    epoch_progress = progress.get_sub_tracker("Epochs")
    epoch_progress.update(total=num_epochs, advance=0)
    for epoch in range(num_epochs):
        _check_control_events(stop_event=stop_event, pause_event=pause_event)
        train_loss, val_loss, train_acc, val_acc = train_epoch(
            model,
            train_loader,
            val_loader,
            optimizer,
            criterion,
            log=log,
            progress=epoch_progress,
            stop_event=stop_event,
            pause_event=pause_event,
        )

        # Update the history
        nhist[epoch, 0] = train_loss
        nhist[epoch, 1] = val_loss
        nhist[epoch, 2] = train_acc
        nhist[epoch, 3] = val_acc

        log.info(
            f"Epoch [{epoch + 1}/{num_epochs}], "
            f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, "
            f"Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}"
        )

        # If we are at 25, 50, or 75% of the epochs, save the model
        if num_epochs > 4:
            if (epoch + 1) % (num_epochs // 4) == 0:
                model_save_path = (
                    output_path / "intermediate_models" / f"model_epoch_{epoch + 1}.pt"
                )
                torch.save(model.state_dict(), model_save_path)
                log.info(f"Model saved at epoch {epoch + 1}")

        epoch_progress.update(total=num_epochs, advance=1)

    # return the trained model and the traning history

    history = pd.DataFrame(
        data=nhist.astype(np.float32),
        columns=["train_loss", "val_loss", "train_acc", "val_acc"],
        index=np.arange(1, num_epochs + 1),
    )

    return model, history


def test_model_bayesian(
    model: nn.Module,
    test_loader: DataLoader[ADNIDataset],
    criterion: nn.Module,
    get_kl_loss: KLLossFn,
    progress: ProgressTracker,
    log: PipelineLogger,
    stop_event: Event | None = None,
    pause_event: Event | None = None,
) -> TestMetrics:
    """
    Tests a Bayesian model on the test dataset with KL-augmented loss.
    """
    model.eval()
    test_loss = 0.0
    correct = 0
    total = 0
    total_samples = 0

    test_progress = progress.get_sub_tracker("Testing Batches")
    test_progress.update(total=len(test_loader), advance=0)
    with torch.no_grad():
        for _, (mri, xls, targets, _) in enumerate(test_loader):
            _check_control_events(stop_event=stop_event, pause_event=pause_event)
            mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
            outputs = model((mri, xls))
            data_loss = cast(torch.Tensor, criterion(outputs, targets))
            batch_size = mri.size(0)
            kl_term = get_kl_loss(model)
            kl_loss = (
                kl_term / batch_size
                if kl_term is not None
                else torch.tensor(0.0, device=outputs.device)
            )
            loss: torch.Tensor = data_loss + kl_loss

            test_loss += loss.item() * batch_size
            total_samples += batch_size

            batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
            correct += batch_correct
            total += batch_total
            test_progress.update(total=len(test_loader), advance=1)

    test_loss = test_loss / total_samples if total_samples > 0 else 0.0
    test_acc = correct / total if total > 0 else 0.0
    return test_loss, test_acc


def train_epoch_bayesian(
    model: nn.Module,
    train_loader: DataLoader[ADNIDataset],
    val_loader: DataLoader[ADNIDataset],
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    get_kl_loss: KLLossFn,
    log: PipelineLogger,
    progress: ProgressTracker,
    stop_event: Event | None = None,
    pause_event: Event | None = None,
) -> TrainMetrics:
    """
    Trains a Bayesian model for one epoch and evaluates on validation data.
    """
    model.train()
    train_loss = 0.0
    train_correct = 0
    train_total = 0
    train_total_samples = 0

    train_progress = progress.get_sub_tracker("Training Batches")
    train_progress.reset()
    train_progress.set_title("Training Batches")
    train_progress.update(total=len(train_loader), advance=0)
    for _, (mri, xls, targets, _) in enumerate(train_loader):
        _check_control_events(stop_event=stop_event, pause_event=pause_event)
        mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
        optimizer.zero_grad()
        outputs = model((mri, xls))
        data_loss = cast(torch.Tensor, criterion(outputs, targets))
        batch_size = mri.size(0)
        kl_term = get_kl_loss(model)
        kl_loss = (
            kl_term / batch_size
            if kl_term is not None
            else torch.tensor(0.0, device=outputs.device)
        )
        loss: torch.Tensor = data_loss + kl_loss
        loss.backward()  # type: ignore[reportUnknownMemberType]
        optimizer.step()
        train_loss += loss.item() * batch_size
        train_total_samples += batch_size
        batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
        train_correct += batch_correct
        train_total += batch_total
        train_progress.update(total=len(train_loader), advance=1)
    train_loss = train_loss / train_total_samples if train_total_samples > 0 else 0.0

    model.eval()
    val_loss = 0.0
    correct = 0
    total = 0
    val_total_samples = 0

    test_progress = progress.get_sub_tracker("Validation Batches")
    test_progress.reset()
    test_progress.set_title("Validation Batches")
    test_progress.update(total=len(val_loader), advance=0)
    with torch.no_grad():
        for _, (mri, xls, targets, _) in enumerate(val_loader):
            _check_control_events(stop_event=stop_event, pause_event=pause_event)
            mri, xls, targets = _move_batch_to_model_device(model, mri, xls, targets)
            outputs = model((mri, xls))
            data_loss = cast(torch.Tensor, criterion(outputs, targets))
            batch_size = mri.size(0)
            kl_term = get_kl_loss(model)
            kl_loss = (
                kl_term / batch_size
                if kl_term is not None
                else torch.tensor(0.0, device=outputs.device)
            )
            vloss: torch.Tensor = data_loss + kl_loss
            val_loss += vloss.item() * batch_size
            val_total_samples += batch_size

            batch_correct, batch_total = _batch_correct_and_total(outputs, targets)
            correct += batch_correct
            total += batch_total
            test_progress.update(total=len(val_loader), advance=1)

    val_loss = val_loss / val_total_samples if val_total_samples > 0 else 0.0
    val_acc = correct / total if total > 0 else 0.0
    train_acc = train_correct / train_total if train_total > 0 else 0.0

    return train_loss, val_loss, train_acc, val_acc


def train_model_bayesian(
    model: nn.Module,
    train_loader: DataLoader[ADNIDataset],
    val_loader: DataLoader[ADNIDataset],
    optimizer: torch.optim.Optimizer,
    criterion: nn.Module,
    num_epochs: int,
    output_path: pl.Path,
    get_kl_loss: KLLossFn,
    progress: ProgressTracker,
    log: PipelineLogger,
    stop_event: Event | None = None,
    pause_event: Event | None = None,
) -> Tuple[nn.Module, pd.DataFrame]:
    """
    Trains a Bayesian model with KL-augmented objective.
    """
    nhist = np.zeros((num_epochs, 4), dtype=np.float32)

    epoch_progress = progress.get_sub_tracker("Epochs")
    epoch_progress.update(total=num_epochs, advance=0)
    for epoch in range(num_epochs):
        _check_control_events(stop_event=stop_event, pause_event=pause_event)
        train_loss, val_loss, train_acc, val_acc = train_epoch_bayesian(
            model,
            train_loader,
            val_loader,
            optimizer,
            criterion,
            get_kl_loss,
            log=log,
            progress=epoch_progress,
            stop_event=stop_event,
            pause_event=pause_event,
        )

        nhist[epoch, 0] = train_loss
        nhist[epoch, 1] = val_loss
        nhist[epoch, 2] = train_acc
        nhist[epoch, 3] = val_acc

        log.info(
            f"Epoch [{epoch + 1}/{num_epochs}], "
            f"Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}, "
            f"Train Acc: {train_acc:.4f}, Val Acc: {val_acc:.4f}"
        )

        if num_epochs > 4:
            if (epoch + 1) % (num_epochs // 4) == 0:
                model_save_path = (
                    output_path / "intermediate_models" / f"model_epoch_{epoch + 1}.pt"
                )
                torch.save(model.state_dict(), model_save_path)
                log.info(f"Model saved at epoch {epoch + 1}")

        epoch_progress.update(total=num_epochs, advance=1)

    history = pd.DataFrame(
        data=nhist.astype(np.float32),
        columns=["train_loss", "val_loss", "train_acc", "val_acc"],
        index=np.arange(1, num_epochs + 1),
    )

    return model, history