Main_3_features.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import packages and functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from From_clusters_to_couples import CreateCouples\n",
    "from Transformations import apply_translation, apply_reduce_to_3_features, apply_randomrotations_3_features\n",
    "from PCT_definition import sample_and_group, farthest_point_sample, index_points, square_distance, Local_op, SA_Layer, StackedAttention\n",
    "from Model_3_features import use_GPU, PairModel1, Branch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the seeds:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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 very similar to Mian.ipynb. The only differences are that 3 features are used, so there is the application of the apply_reduce_to_3_features() function at the beginning of the training cycle that extracts the first 3 features from the data. Also the apply_randomrotations_3_features() function is used, since apply_randomrotations() requires at least 6 features."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_couples = CreateCouples(train)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Identify positive and negative couples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "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\\\\Main3_epoch63.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 = \"three_features\",\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",
    "        # Reduce to three features\n",
    "        frags_a= apply_reduce_to_3_features(frags_a)\n",
    "        frags_b= apply_reduce_to_3_features(frags_b)\n",
    "\n",
    "        # Random rotations\n",
    "        frags_a = apply_randomrotations_3_features(frags_a)\n",
    "        frags_b = apply_randomrotations_3_features(frags_b)\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",
    "            # Reduce to three features\n",
    "            val_frags_a= apply_reduce_to_3_features(val_frags_a)\n",
    "            val_frags_b= apply_reduce_to_3_features(val_frags_b)\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\\\\Main3_epoch63.pt')\n",
    "model.load_state_dict(W_stored)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the Test dataloader:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loader_basic = DataLoader(test_randomized, batch_size=16) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set hyperparameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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",
    "        # Reduce to three features\n",
    "        val_frags_a= apply_reduce_to_3_features(val_frags_a)\n",
    "        val_frags_b= apply_reduce_to_3_features(val_frags_b)\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
}