run_experiment.py

import argparse

# from logger import CheckpointSaver
import logging
import os
import random

import numpy as np
import torch
from torch import nn
from torchmetrics import Accuracy, F1Score
from tqdm import tqdm

import wandb

logging.getLogger().setLevel(logging.INFO)

from my_utilities.config_reader import read_config
from my_utilities.data_provider_NN import get_data
from my_utilities.stopper import EarlyStopper


# Dataloader randomness
# if num_workers >1 https://pytorch.org/docs/stable/notes/randomness.html#dataloader
def set_seed(seed: int = 42) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    # when running on the CuDNN backend, two further options must be set
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    # set a fixed value for the hash seed
    os.environ["PYTHONHASHSEED"] = str(seed)
    print(f"Random seed set as {seed}")


def train_fn(model, criterion, train_data_loader, optimizer, epoch, device="cuda"):
    model.train()
    accuracy = Accuracy(task="binary").to(device)
    f1 = F1Score(task="binary").to(device)
    fin_loss = 0.0
    fin_accuracy = 0.0
    fin_f1 = 0.0
    tk = tqdm(train_data_loader, desc="Epoch" + " [TRAIN] " + str(epoch + 1))
    for t, data in enumerate(tk):
        inputs, labels, _ = data
        inputs, labels = inputs.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)

        labels = labels.float()

        train_loss = criterion(outputs, labels)
        fin_accuracy += accuracy(outputs, labels)
        fin_f1 += f1(outputs, labels)
        train_loss.backward()
        optimizer.step()

        fin_loss += train_loss.item()
        tk.set_postfix(
            {
                "loss": "%.6f" % float(fin_loss / (t + 1)),
                "LR": optimizer.param_groups[0]["lr"],
            }
        )
    return (
        fin_loss / len(train_data_loader),
        fin_accuracy / len(train_data_loader),
        fin_f1 / len(train_data_loader),
        optimizer.param_groups[0]["lr"],
    )


def eval_fn(model, criterion, eval_data_loader, epoch, device="cuda"):
    model.eval()
    accuracy = Accuracy(task="binary").to(device)
    f1 = F1Score(task="binary").to(device)
    fin_loss = 0.0
    fin_accuracy = 0.0
    fin_f1 = 0.0
    tk = tqdm(eval_data_loader, desc="Epoch" + " [VALID] " + str(epoch + 1))

    with torch.no_grad():
        for t, data in enumerate(tk):
            inputs, labels, _ = data
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            labels = labels.float()
            eval_loss = criterion(outputs, labels)
            fin_loss += eval_loss.item()
            fin_accuracy += accuracy(outputs, labels)
            fin_f1 += f1(outputs, labels)
            tk.set_postfix({"loss": "%.6f" % float(fin_loss / (t + 1))})
        return (
            fin_loss / len(eval_data_loader),
            fin_accuracy / len(eval_data_loader),
            fin_f1 / len(eval_data_loader),
        )


class CheckpointSaver:
    def __init__(
        self,
        experiment,
        dirpath,
        save_all=False,
        decreasing=True,
        top_n=999999999,
    ):
        """
        dirpath: Directory path where to store all model weights
        decreasing: If decreasing is `True`, then lower metric is better
        top_n: Total number of models to track based on validation metric value
        """
        if not os.path.exists(dirpath):
            os.makedirs(dirpath)
        self.dirpath = dirpath
        self.top_n = top_n
        self.decreasing = decreasing
        self.best_metric_val = np.Inf if decreasing else -np.Inf
        self.experiment = experiment
        self.save_all = save_all

    def __call__(self, model, epoch, metric_to_save, **kwargs):
        model_name = f"{self.experiment}-ep-{epoch}.pt"
        model_path = os.path.join(self.dirpath, model_name)
        if self.save_all:
            model_scripted = torch.jit.script(model)
            model_scripted.save(model_path)
            self.log_artifact_custom(model_name, model_path, epoch, **kwargs)
        else:
            save = (
                metric_to_save < self.best_metric_val
                if self.decreasing
                else metric_to_save > self.best_metric_val
            )
            if save:
                logging.info(
                    f"Current metric value better than {metric_to_save} better than best {self.best_metric_val}, saving model at {model_path}, & logging model weights to W&B."
                )
                self.best_metric_val = metric_to_save
                # TODO: ID1: change saving model approach:
                #            create class that initiate model or basically save model definition
                #            model definition can be in Yaml file or ??
                #            according to structure in config
                model_scripted = torch.jit.script(model)
                model_scripted.save(model_path)

                self.log_artifact_custom(model_name, model_path, epoch, **kwargs)

    def log_artifact_custom(self, filename, model_path, epoch, **kwargs):
        metadata = {"epoch": epoch}
        for k, v in kwargs.items():
            metadata[k] = v
        artifact = wandb.Artifact(filename, type="model", metadata=metadata)

        # TODO: save on the fly
        artifact.add_file(model_path)
        wandb.run.log_artifact(artifact)


def run_train(config=None):
    """
    Function is taking care of:
        1. Initial data reading
        2. Data loaders preparation
        3. Formation of model structure
        4. Checkpoint saver init
        5. Iteration over epochs
        6. Calling training and evaluation
    """
    DATASET_PATH = "/exploiting_model_multiplicity/data/"
    # Initialize a new wandb run
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    print(f"Using {device} device")

    with wandb.init(config=config) as run:
        config = wandb.config
        config.experiment_name = "Expr" + "-" + run.name.split("-")[2]
        run.name = "Expr" + "-" + run.name.split("-")[2]
        # If called by wandb.agent, as below,
        # this config will be set by Sweep Controller
        set_seed(config.seed)

        script_dir = os.path.dirname(os.path.abspath(__file__))

        train_dir = os.path.join(script_dir, "data", config.dataset, "train")
        test_dir = os.path.join(script_dir, "data", config.dataset, "test")
        print("---- PATH -----")
        print(train_dir)

        trainloader = get_data(train_dir, batch_size=config.batch_size)
        testloader = get_data(test_dir, batch_size=config.batch_size)
        # model
        input_size = testloader.dataset.X.shape[1]
        hidden_sizes = [40, 20]
        output_size = 1
        # TODO: pass dictionary style the model architecture to this function
        model = nn.Sequential(
            nn.Linear(input_size, hidden_sizes[0]),
            nn.ReLU(),
            nn.Linear(hidden_sizes[0], hidden_sizes[1]),
            nn.ReLU(),
            nn.Linear(hidden_sizes[1], output_size),
            nn.Sigmoid(),
        )

        model = model.to(device)
        criterion = nn.BCELoss()
        # optimizer
        # TODO: pass optimiser to this function
        # TODO: test if optimiser can be passed from config
        optimizer = torch.optim.SGD(model.parameters(), lr=config.lr)
        # source: https://www.kaggle.com/code/isbhargav/guide-to-pytorch-learning-rate-scheduling
        if config.lr_decay:
            lambda1 = lambda epoch: 0.65**epoch
            scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1)

        # checkpoint saver
        # TODO: pass param top_n
        checkpoint_saver = CheckpointSaver(
            config.experiment_name,
            dirpath="./model_weights",
            save_all=True,
            decreasing=True,
            top_n=5,
        )
        early_stopper = EarlyStopper(
            patience=5,
            min_delta=0,
            decreasing=False,
        )
        for epoch in range(config.epochs):
            print(f"epoch started: {epoch}")
            avg_loss_train, avg_accuracy_train, avg_f1_train, lr = train_fn(
                model, criterion, trainloader, optimizer, epoch, device=device
            )
            avg_loss_eval, avg_accuracy_eval, avg_f1_eval = eval_fn(
                model, criterion, testloader, epoch, device=device
            )

            checkpoint_saver(
                model,
                epoch,
                metric_to_save=avg_loss_eval,
                train_loss=avg_loss_train,
                eval_loss=avg_loss_eval,
                train_acc=avg_accuracy_train,
                train_f1=avg_f1_train,
                eval_acc=avg_accuracy_eval,
                eval_f1=avg_f1_eval,
                dataset=config.dataset,
                seed=config.seed,
                batch_size=config.batch_size,
                lr=lr,
            )

            wandb.run.log(
                {
                    "epoch": epoch,
                    "train loss": avg_loss_train,
                    "eval loss": avg_loss_eval,
                    "train acc": avg_accuracy_train,
                    "train f1": avg_f1_train,
                    "eval acc": avg_accuracy_eval,
                    "eval f1": avg_f1_eval,
                    "lr": lr,
                }
            )
            print(f"LEARNING RATE BEFORE {lr:.5f}")
            print("------------------------------")
            if config.lr_decay:
                scheduler.step()
            print(
                f"EP = {epoch} | TR_LOSS = {avg_loss_train:.2f} | EVL_LOSS = {avg_loss_eval:.2f}"
            )
            print(
                f"TR_ACC = {avg_accuracy_train:.2f} | EVL_ACC = {avg_accuracy_eval:.2f}"
            )
            print(f"TR_F1 = {avg_f1_train:.2f} | EVL_F1 = {avg_f1_eval:.2f}")
            print("------------------------------")
            print(f"LEARNING RATE AFTER {lr:.5f}")
            if early_stopper.early_stop(avg_accuracy_eval):
                break


def init_experiment(args):
    if args.config and args.sweep:
        print(
            "Error: You cannot provide both a configuration file and a seep_id. Please provide only one."
        )
        return -1
    if wandb.login(key=""):
        if args.config:
            config = read_config(args.config)
            # Initialize sweep by passing in config.
            # (Optional) Provide a name of the project.
            sweep_id = wandb.sweep(
                sweep=config,
                project=config.get("project"),
            )
            print(f"sweep {sweep_id} created")
            wandb.agent(
                sweep_id,
                function=run_train,
            )

        else:
            sweep_id = args.sweep
            print("Starting agent with seep_id")
            sweep_id = ""
            wandb.agent(
                sweep_id,
                function=run_train,
                count=args.count,
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process data with input parameters")
    parser.add_argument("--config", help="Config name", type=str, default=None)
    parser.add_argument("--sweep", help="Sweep ID", type=str, default=None)
    parser.add_argument("--count", help="Agent runs", type=int, default=10)
    args = parser.parse_args()
    init_experiment(args)