Main_Not_Augmentation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages and functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "import pandas as pd\n",
    "import random\n",
    "import pickle\n",
    "import datetime\n",
    "import os\n",
    "import torch.nn.functional as F\n",
    "from scipy.stats import special_ortho_group\n",
    "from scipy.spatial.transform import Rotation as R\n",
    "import torch.nn.functional as F\n",
    "import wandb\n",
    "import itertools\n",
    "from sklearn.metrics import f1_score, roc_curve, auc, roc_auc_score, confusion_matrix\n",
    "from torch.nn.functional import softmax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from From_clusters_to_couples import CreateCouples\n",
    "from Transformations import apply_translation, apply_randomrotations\n",
    "from PCT_definition import sample_and_group, farthest_point_sample, index_points, square_distance, Local_op, SA_Layer, StackedAttention\n",
    "from Model_7_features import use_GPU, PairModel1, Branch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the seeds:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seed=999\n",
    "os.environ[\"PL_GLOBAL_SEED\"] = str(seed)\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed_all(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code is the same as Main.ipynb, the only difference is that apply_randomrotations() is not applied in the train loop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = torch.load(\"C:\\\\Users\\\\Alessandro\\\\Desktop\\\\Tesi\\\\pair_dataset\\\\dataset_1024_AB\\\\train_pair_dataset_REG.pt\")\n",
    "\n",
    "with open('C:\\\\Users\\\\Alessandro\\\\Desktop\\\\Tesi\\\\PairModel\\\\val_subset_random.pkl', 'rb') as file:\n",
    "    val_randomized = pickle.load(file)\n",
    "\n",
    "with open('C:\\\\Users\\\\Alessandro\\\\Desktop\\\\Tesi\\\\PairModel\\\\test_subset_random.pkl', 'rb') as file:\n",
    "    test_randomized = pickle.load(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Unroll the pairs from the clusters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_couples = CreateCouples(train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Identify positive and negative couples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_couples_0 = [item for item in train_couples if item[2] == 0]\n",
    "train_couples_1 = [item for item in train_couples if item[2] == 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the validation dataloader:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_loader_basic = DataLoader(val_randomized, batch_size=16) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Activate GPU:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device=use_GPU()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = PairModel1().to(device)\n",
    "model.double()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Only when using pre-trained weights:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W_stored = torch.load(r'C:\\\\Users\\\\Alessandro\\\\Desktop\\\\Weights\\\\Main_not_aug_epoch98.pt')\n",
    "model.load_state_dict(W_stored)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set hyperparameters and prepare Weights & Biases:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.init(\n",
    "      project=\"train\", \n",
    "      notes = \"No Data Augmentation\",\n",
    "      config={\n",
    "      \"learning_rate\": 0.00005,\n",
    "      \"architecture\": \"Model2_mod\",\n",
    "      \"epochs\": 20,\n",
    "      \"weight_decay\": 0.0001,\n",
    "      \"W_crossentropy\":1,\n",
    "      \"couples_per_epoch\": 10000,\n",
    "      \"seed\": seed # the seed defined at the beginning\n",
    "      })\n",
    "      \n",
    "config = wandb.config\n",
    "criterion = nn.CrossEntropyLoss().to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate, weight_decay=config.weight_decay)\n",
    "num_epochs = config.epochs\n",
    "best_val_accuracy = 0.0 \n",
    "\n",
    "\n",
    "# Sets the path where the model weights will be stored.\n",
    "# make sure to enter the right destination folder, otherwise the training cycle will stop by not finding the folder\n",
    "checkpoint_dir = r'C:\\\\Users\\\\Alessandro\\\\Desktop\\\\Tesi\\\\PairModel\\\\Check_points'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training cycle:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint_interval = 1\n",
    "epoch_number = 0  \n",
    "\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train() \n",
    "\n",
    "    total_loss = 0.0\n",
    "    correct_predictions = 0\n",
    "    total_samples = 0\n",
    "    y_true = []\n",
    "    y_pred = []\n",
    "    \n",
    "    random.shuffle(train_couples_0)\n",
    "    random.shuffle(train_couples_1)\n",
    "\n",
    "    # 5000 positive pairs and 5000 negative pairs are used\n",
    "    balanced_train_list = []\n",
    "    couples_per_epoch = config.couples_per_epoch\n",
    "    \n",
    "    sampled_couples_0 = random.sample(train_couples_0, couples_per_epoch // 2)\n",
    "    sampled_couples_1 = random.sample(train_couples_1, couples_per_epoch // 2)\n",
    "\n",
    "    balanced_train_list.extend(sampled_couples_0)\n",
    "    balanced_train_list.extend(sampled_couples_1)\n",
    "        \n",
    "    random.shuffle(balanced_train_list)\n",
    "    \n",
    "    train_loader = DataLoader(balanced_train_list, batch_size=16,shuffle=True) \n",
    "    \n",
    "    progress_bar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs}', leave=False)\n",
    "\n",
    "    ###########\n",
    "    ## Train ##\n",
    "    ###########\n",
    "\n",
    "\n",
    "    for batch_data in progress_bar:\n",
    "\n",
    "        optimizer.zero_grad() \n",
    "        frags_a, frags_b, labels = batch_data\n",
    "        \n",
    "        # No random rotations applied\n",
    "\n",
    "        # Translations in the origin\n",
    "        frags_a = apply_translation(frags_a)\n",
    "        frags_b = apply_translation(frags_b)\n",
    "\n",
    "        frags_a = frags_a.double().to(device)\n",
    "        frags_b = frags_b.double().to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        # Model output\n",
    "        outputs = model(frags_a, frags_b)\n",
    "        loss = criterion(outputs, labels)\n",
    "\n",
    "        # Optimization\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        _, predicted = torch.max(outputs.data, 1)\n",
    "        total_samples += labels.size(0)\n",
    "        correct_predictions += (predicted == labels).sum().item()\n",
    "\n",
    "        y_true.extend(labels.cpu().numpy())\n",
    "        y_pred.extend(predicted.cpu().numpy())\n",
    "        progress_bar.set_postfix({'Loss': loss.item(), 'Accuracy': correct_predictions / total_samples})\n",
    "        \n",
    "    # Metrics computation\n",
    "    f1 = f1_score(y_true, y_pred, average='weighted')\n",
    "    accuracy = correct_predictions / total_samples\n",
    "    train_loss = total_loss/len(train_loader)\n",
    "\n",
    "    metrics_train = {\"train_loss\": train_loss, \n",
    "                       \"accuracy\": accuracy,\n",
    "                       \"f1\":f1}\n",
    "    \n",
    "    # Upload results to wandb\n",
    "    wandb.log(metrics_train) \n",
    "\n",
    "\n",
    "\n",
    "    ###############\n",
    "    ## Inference ##\n",
    "    ###############\n",
    "\n",
    "    model.eval()  \n",
    "    \n",
    "    val_loss = 0.0\n",
    "    val_correct_predictions = 0\n",
    "    val_total_samples = 0\n",
    "    y_true_val = []\n",
    "    y_pred_val = []\n",
    "    with torch.no_grad():\n",
    "        for val_batch in val_loader_basic:\n",
    "            val_frags_a, val_frags_b, val_labels = val_batch\n",
    "            \n",
    "            # Translations in the origin\n",
    "            val_frags_a = apply_translation(val_frags_a)\n",
    "            val_frags_b = apply_translation(val_frags_b)\n",
    "\n",
    "            val_frags_a = val_frags_a.double().to(device)\n",
    "            val_frags_b = val_frags_b.double().to(device)\n",
    "            val_labels = val_labels.to(device)\n",
    "\n",
    "            # Model output\n",
    "            val_outputs = model(val_frags_a, val_frags_b)\n",
    "\n",
    "            val_loss += criterion(val_outputs, val_labels).item()\n",
    "            _, val_predicted = torch.max(val_outputs.data, 1)\n",
    "            val_total_samples += val_labels.size(0)\n",
    "            val_correct_predictions += (val_predicted == val_labels).sum().item()\n",
    "            y_true_val.extend(val_labels.cpu().numpy())\n",
    "            y_pred_val.extend(val_predicted.cpu().numpy())\n",
    "\n",
    "    # Metrics computation\n",
    "    f1_val = f1_score(y_true_val, y_pred_val, average='weighted')        \n",
    "    val_accuracy = val_correct_predictions / val_total_samples\n",
    "    val_loss /= len(val_loader_basic)\n",
    "    val_metrics = {\"val_loss\": val_loss, \n",
    "                       \"val_accuracy\": val_accuracy,\n",
    "                       \"f1_val\": f1_val}\n",
    "    \n",
    "    # Upload results to wandb\n",
    "    wandb.log(val_metrics)\n",
    "\n",
    "    print(f'Epoch [{epoch+1}/{num_epochs}], Training Loss: {train_loss:.4f}, Training Accuracy: {accuracy:.4f}, F1_train:{f1:.4f} ',\n",
    "    f'Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.4f}, F1_val : {f1_val:.4f}')\n",
    "\n",
    "\n",
    "    # Store the results    \n",
    "    current_time = datetime.datetime.now()\n",
    "    checkpoint_name = f\"{current_time.strftime('%m%d_%H%M%S')}_{epoch + 1}.pt\"    \n",
    "    # make sure to enter the right destination folder, otherwise the training cycle will stop by not finding the folder\n",
    "    checkpoint_path = os.path.join(checkpoint_dir, checkpoint_name) \n",
    "    torch.save(model.state_dict(), checkpoint_path)\n",
    "    wandb.run.log_artifact(checkpoint_path,name=str(epoch+1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference time:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pre-trained weight loading:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "W_stored = torch.load(r'C:\\\\Users\\\\Alessandro\\\\Desktop\\\\Weights\\\\Main_not_aug_epoch98.pt')\n",
    "model.load_state_dict(W_stored)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the Test dataloader:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader_basic = DataLoader(test_randomized, batch_size=16) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set hyperparameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss().to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.00005, weight_decay=0.0001)\n",
    "num_epochs = 1\n",
    "best_val_accuracy = 0.0 \n",
    "# Sets the path where the model weights will be stored.\n",
    "checkpoint_dir = r'C:\\\\Users\\\\Alessandro\\\\Desktop\\\\Tesi\\\\PairModel\\\\Check_points'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inference cycle:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()  \n",
    "progress_bar_val = tqdm(test_loader_basic, desc=f'Validation', leave=False)    \n",
    "val_loss = 0.0\n",
    "val_correct_predictions = 0\n",
    "val_total_samples = 0\n",
    "\n",
    "y_true_val = []\n",
    "y_pred_val = []\n",
    "y_scores_val = []\n",
    "\n",
    "with torch.no_grad():\n",
    "    for val_batch in progress_bar_val:\n",
    "        val_frags_a, val_frags_b, val_labels = val_batch\n",
    "        \n",
    "        # Translations in the origin\n",
    "        val_frags_a = apply_translation(val_frags_a)\n",
    "        val_frags_b = apply_translation(val_frags_b)\n",
    "        \n",
    "        val_frags_a = val_frags_a.double().to(device)\n",
    "        val_frags_b = val_frags_b.double().to(device)\n",
    "\n",
    "        val_labels = val_labels.to(device)\n",
    "            \n",
    "        val_outputs = model(val_frags_a, val_frags_b)\n",
    "        \n",
    "        val_outputs_probs = softmax(val_outputs, dim=1)\n",
    "        val_loss += criterion(val_outputs, val_labels).item()\n",
    "        _, val_predicted = torch.max(val_outputs.data, 1)\n",
    "        val_total_samples += val_labels.size(0)\n",
    "        val_correct_predictions += (val_predicted == val_labels).sum().item()\n",
    "\n",
    "        y_true_val.extend(val_labels.cpu().numpy())\n",
    "        y_pred_val.extend(val_predicted.cpu().numpy())\n",
    "        y_scores_val.extend(val_outputs_probs.cpu().numpy()[:, 1])  \n",
    "        \n",
    "# ROC curve\n",
    "fpr, tpr, thresholds = roc_curve(y_true_val, y_scores_val)\n",
    "\n",
    "# AUC\n",
    "roc_auc = roc_auc_score(y_true_val, y_scores_val)\n",
    "\n",
    "# F1-score, Accuracy, Val_loss\n",
    "f1_val = f1_score(y_true_val, y_pred_val, average='weighted') \n",
    "val_accuracy = val_correct_predictions / val_total_samples\n",
    "val_loss /= len(test_loader_basic)\n",
    "\n",
    "# Print\n",
    "print(\"Test Accuracy:\", val_accuracy)\n",
    "print(\"Test F1 Score:\", f1_val)\n",
    "print(\"Test AUC:\", roc_auc)\n",
    "print(\"Test Loss\", val_loss)\n",
    "\n",
    "conf_matrix = confusion_matrix(y_true_val, y_pred_val)\n",
    "\n",
    "# Print the confusion matrix\n",
    "print(\"Confusion Matrix:\")\n",
    "print(conf_matrix)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}