training-testing/training-transformerGO.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "undefined-collar",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Ieremie\\Desktop\\TransformerGO\n",
      "C:\\Users\\Ieremie\\Desktop\\TransformerGO\\protein-ppi-encoding-module\n",
      "importing Jupyter notebook from transformerGO.ipynb\n",
      "importing Jupyter notebook from harvard_transformer.ipynb\n",
      "C:\\Users\\Ieremie\\Desktop\\TransformerGO\n",
      "C:\\Users\\Ieremie\\Desktop\\TransformerGO\\datasets\n",
      "importing Jupyter notebook from dataset_manip.ipynb\n",
      "C:\\Users\\Ieremie\\Desktop\\TransformerGO\n",
      "C:\\Users\\Ieremie\\Desktop\\TransformerGO\\training-testing\n",
      "importing Jupyter notebook from training_helper.ipynb\n",
      "C:\\Users\\Ieremie\\Desktop\\TransformerGO\n",
      "Device available:  cuda   GeForce GTX 960M\n"
     ]
    }
   ],
   "source": [
    "import import_ipynb\n",
    "\n",
    "#import model\n",
    "%cd ..\n",
    "%cd \"protein-ppi-encoding-module\"\n",
    "from transformerGO import *\n",
    "\n",
    "#import dataset functions\n",
    "%cd ..\n",
    "%cd \"datasets\"\n",
    "from dataset_manip import *\n",
    "%cd .. \n",
    "\n",
    "%cd \"training-testing\"\n",
    "from training_helper import * \n",
    "%cd ..\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.utils.data as data\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "import torch.optim as optim\n",
    "\n",
    "import time\n",
    "import pickle\n",
    "\n",
    "import random\n",
    "from random import shuffle\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from torchsummary import summary\n",
    "\n",
    "from sklearn import metrics\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from tqdm import tqdm as prog_bar\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(\"Device available: \", device, \" \", torch.cuda.get_device_name(0))\n",
    "\n",
    "#for np array from nested sequences\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\", category=np.VisibleDeprecationWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4e443cb3-4898-48f8-880f-82a36ce6b0f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def helper_collate(batch):\n",
    "    MAX_LEN_SEQ = get_max_len_seq(batch)\n",
    "    return transformerGO_collate_fn(batch, MAX_LEN_SEQ, EMB_DIM, pytorch_pad = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "active-navigator",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rejected interactions where at least one protein has no annotation:  6573\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 500]:  0\n",
      "Number of interactions: 262570\n",
      "Rejected interactions where at least one protein has no annotation:  38444\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 500]:  0\n",
      "Number of interactions: 230699\n",
      "Rejected interactions where at least one protein has no annotation:  1657\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 500]:  0\n",
      "Number of interactions: 65628\n",
      "Rejected interactions where at least one protein has no annotation:  9620\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 500]:  0\n",
      "Number of interactions: 57665\n",
      "Rejected interactions where at least one protein has no annotation:  2018\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 500]:  0\n",
      "Number of interactions: 82088\n",
      "Rejected interactions where at least one protein has no annotation:  11801\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 500]:  0\n",
      "Number of interactions: 72305\n"
     ]
    }
   ],
   "source": [
    "## STRINGDB BENCHMARK DATASET ##\n",
    "EMB_DIM = 64\n",
    "organism = 9606\n",
    "data_path = 'datasets/onto2vec-datasets-string/data'\n",
    "\n",
    "go_embed_pth = f\"term-encoding-module/emb/go-terms-{EMB_DIM}.emd\"\n",
    "go_id_dict_pth = \"term-encoding-module/go_id_dict\"\n",
    "protein_go_anno_pth = data_path +\"/goa_human.gaf.gz\"\n",
    "alias_path = data_path + f'/{organism}.protein.aliases.v11.0.txt.gz'\n",
    "\n",
    "neg_path_train = data_path + f'/train/{organism}.no-mirror.negative_interactions.txt'\n",
    "poz_path_train = data_path + f'/train/{organism}.no-mirror.protein.links.v11.0.txt'\n",
    "neg_path_valid = data_path + f'/valid/{organism}.no-mirror.negative_interactions.txt'\n",
    "poz_path_valid = data_path + f'/valid/{organism}.no-mirror.protein.links.v11.0.txt'\n",
    "neg_path_test = data_path + f'/test/{organism}.no-mirror.negative_interactions.txt'\n",
    "poz_path_test = data_path + f'/test/{organism}.no-mirror.protein.links.v11.0.txt'\n",
    "\n",
    "#the splitting for StringDB is already done, just load the data\n",
    "train_set = get_dataset_split_stringDB(poz_path_train, neg_path_train, protein_go_anno_pth, go_id_dict_pth, go_embed_pth, shuffle,alias_path, ratio = [1, 0, 0], stringDB = True)[0]\n",
    "valid_set = get_dataset_split_stringDB(poz_path_valid, neg_path_valid, protein_go_anno_pth, go_id_dict_pth, go_embed_pth, shuffle,alias_path, ratio = [1, 0, 0], stringDB = True)[0]\n",
    "test_set = get_dataset_split_stringDB(poz_path_test, neg_path_test, protein_go_anno_pth, go_id_dict_pth, go_embed_pth, shuffle,alias_path, ratio = [1, 0, 0], stringDB = True)[0]\n",
    "\n",
    "params = {'batch_size': 32,'collate_fn': helper_collate}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "95bbdde5-801c-477f-ad35-3e922a70abb2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rejected interactions where at least one protein has no annotation:  9424\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 5000]:  0\n",
      "Number of interactions: 243560\n",
      "Rejected interactions where at least one protein has no annotation:  10385\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 5000]:  0\n",
      "Number of interactions: 242599\n"
     ]
    }
   ],
   "source": [
    "#TRANSFORMERGO DATASET#\n",
    "organism = 9606\n",
    "EMB_DIM = 64\n",
    "data_path = 'datasets/transformerGO-dataset/'\n",
    "go_embed_pth = data_path + f\"go-terms/emb/go-terms-{EMB_DIM}.emd\"\n",
    "go_id_dict_pth = data_path + \"go-terms/go_id_dict\"\n",
    "protein_go_anno_pth = data_path +\"stringDB-files/goa_human.gaf.gz\"\n",
    "alias_path = data_path + f'stringDB-files/{organism}.protein.aliases.v11.5.txt.gz'\n",
    "\n",
    "neg_path = data_path + f'interaction-datasets/{organism}.protein.negative.v11.5.txt'\n",
    "poz_path= data_path + f'interaction-datasets/{organism}.protein.links.v11.5.txt'\n",
    "\n",
    "#choosing 16% data for validation to be similar with opa2vec where validation is 20% of training data (80%)\n",
    "intr_set_size_filter = [0,5000]\n",
    "go_filter = 'ALL'\n",
    "train_set, valid_set, test_set, _ = get_dataset_split_stringDB(poz_path, neg_path, protein_go_anno_pth, go_id_dict_pth, go_embed_pth,\\\n",
    "                                        shuffle, alias_path, ratio = [0.64, 0.16, 0.2], go_filter = go_filter, intr_set_size_filter = intr_set_size_filter, max_intr_size = None)\n",
    "\n",
    "params = {'batch_size': 32,'collate_fn': helper_collate}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "visible-cancer",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<string>:153: DtypeWarning: Columns (16) have mixed types.Specify dtype option on import or set low_memory=False.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rejected interactions where at least one protein has no annotation:  0\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 500]:  0\n",
      "Number of interactions: 1288\n",
      "Rejected interactions where at least one protein has no annotation:  0\n",
      "Rejected interactions where go_filter=ALL and intr_set_size_filter=[0, 500]:  0\n",
      "Number of interactions: 1288\n"
     ]
    }
   ],
   "source": [
    "## JAINS TCSS DATASETS ##\n",
    "EMB_DIM = 64\n",
    "neg_path = \"datasets/jains-TCSS-datasets/human_data/iea-/negatives.human.f\"\n",
    "poz_path = \"datasets/jains-TCSS-datasets/human_data/iea-/positives.human.f\"\n",
    "\n",
    "go_embed_pth = \"term-encoding-module/emb/go-terms-64.emd\"\n",
    "go_id_dict_pth = \"term-encoding-module/go_id_dict\"\n",
    "protein_go_anno_pth = \"datasets/jains-TCSS-datasets/human_data/gene_association.goa_human\"\n",
    "\n",
    "train_set, valid_set, test_set, all_dataset = get_dataset_split_stringDB(poz_path, neg_path, protein_go_anno_pth, go_id_dict_pth, go_embed_pth, shuffle, \"\", ratio = [0.8, 0.2, 0], stringDB = False)\n",
    "params = {'batch_size': 16,'collate_fn': helper_collate}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d14ad123-8a0c-4146-8374-c4ba4d83bb95",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train set:  2060 \n",
      " Valid set:  515 \n",
      " Test set:  1 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Train set: \", len(train_set), '\\n', \"Valid set: \", len(valid_set), '\\n', \"Test set: \", len(test_set), '\\n')\n",
    "train_grt = data.DataLoader(train_set, **params, shuffle = True)\n",
    "val_grt = data.DataLoader(valid_set, **params, shuffle = True)\n",
    "test_grt = data.DataLoader(test_set, **params, shuffle = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "legendary-perception",
   "metadata": {},
   "outputs": [],
   "source": [
    "#TRANSFORMERGO MODEL #\n",
    "MODEL_SIZE = EMB_DIM\n",
    "NR_HEADS = 8\n",
    "NR_LAYERS = 3\n",
    "DROPOUT = 0.2\n",
    "SIZE_FF = 4 * MODEL_SIZE\n",
    "LR = 0.0001\n",
    "\n",
    "model = TransformerGO_Scratch(MODEL_SIZE, NR_HEADS, NR_LAYERS, SIZE_FF, DROPOUT)\n",
    "#model = TransformerGO(MODEL_SIZE, NR_HEADS, NR_LAYERS, SIZE_FF, DROPOUT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3afe3e11-a818-41be-a2db-5744f9400981",
   "metadata": {},
   "outputs": [],
   "source": [
    "#SIMPLE FEED FORWARD NEURAL NETORK#\n",
    "MODEL_SIZE = EMB_DIM\n",
    "DROPOUT = 0.2\n",
    "LR = 0.0001\n",
    "F_1 = 200\n",
    "F_2 = 400\n",
    "F_3 = 200\n",
    "F_4 = 1\n",
    "\n",
    "model =  GO_Sum_NN(MODEL_SIZE, F_1, F_2, F_3, F_4, DROPOUT)\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "92a3a696-8d73-4f89-8161-769a1f09e37c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "350529\n"
     ]
    }
   ],
   "source": [
    "model = model.to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=LR) \n",
    "criterion = nn.BCEWithLogitsLoss().to(device)\n",
    "\n",
    "pytorch_total_params = sum(p.numel() for p in model.parameters())\n",
    "print(pytorch_total_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "vanilla-mexican",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion,  torch_vers = False):\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    epoch_acc = 0\n",
    "    model.train()\n",
    "    \n",
    "    pred = []\n",
    "    lab = []\n",
    "    for batch in prog_bar(iterator):\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        #padded pairs: tensor of shape N * 2(protein pair) * L(longest seq) * Emb dim\n",
    "        padded_pairs = batch[0].to(device)\n",
    "        labels = batch[1].to(device)\n",
    "        mask = batch[2].to(device)\n",
    "        \n",
    "        #split data into protA and protB\n",
    "        gosetA_batch = padded_pairs[:,0]\n",
    "        gosetB_batch = padded_pairs[:,1]\n",
    "        \n",
    "        #permute the data to fit the pytorch transformer\n",
    "        if torch_vers:\n",
    "            gosetA_batch = gosetA_batch.permute(1,0,2)\n",
    "            gosetB_batch = gosetB_batch.permute(1,0,2)\n",
    "        \n",
    "        predictions = model(gosetA_batch, gosetB_batch, mask[:,0], mask[:,1]).squeeze(1)\n",
    "        loss = criterion(predictions, labels)\n",
    "        acc = binary_accuracy(predictions, labels)\n",
    "        \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        epoch_loss += loss.item()\n",
    "        epoch_acc += acc.item()\n",
    "        \n",
    "        pred = pred + list(predictions.cpu().data.numpy())\n",
    "        lab = lab + list(labels.cpu().data.numpy())\n",
    " \n",
    "    return epoch_loss / len(iterator), epoch_acc / len(iterator), roc_auc_score(lab,pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "potential-practice",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, iterator, criterion, torch_vers = False):\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    epoch_acc = 0\n",
    "    pred = []\n",
    "    lab = []\n",
    "    \n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in iterator:\n",
    "            \n",
    "            #padded pairs: tensor of shape N * 2(protein pair) * L(longest seq) * Emb dim\n",
    "            padded_pairs = batch[0].to(device)\n",
    "            labels = batch[1].to(device)\n",
    "            mask = batch[2].to(device)\n",
    "        \n",
    "            #split data into protA and protB\n",
    "            gosetA_batch = padded_pairs[:,0]\n",
    "            gosetB_batch = padded_pairs[:,1]\n",
    "            \n",
    "            #permute the data to fit the pytorch transformer\n",
    "            if torch_vers:\n",
    "                gosetA_batch = gosetA_batch.permute(1,0,2)\n",
    "                gosetB_batch = gosetB_batch.permute(1,0,2)\n",
    "        \n",
    "            predictions = model(gosetA_batch, gosetB_batch, mask[:,0], mask[:,1]).squeeze(1)\n",
    "            loss = criterion(predictions, labels)\n",
    "            acc = binary_accuracy(predictions, labels)\n",
    "            \n",
    "            epoch_loss += loss.item()\n",
    "            epoch_acc += acc.item()\n",
    "            \n",
    "            pred = pred + list(predictions.cpu().data.numpy())\n",
    "            lab = lab + list(labels.cpu().data.numpy())\n",
    "\n",
    "    return epoch_loss / len(iterator), epoch_acc / len(iterator), roc_auc_score(lab,pred), lab, pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "7b79601f-41f3-433a-b46a-7aaacbc0fe89",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_cosine(iterator, criterion, torch_vers = False):\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    epoch_acc = 0\n",
    "    pred = []\n",
    "    lab = []\n",
    "    \n",
    "    for batch in prog_bar(iterator):\n",
    "            \n",
    "        #padded pairs: tensor of shape N * 2(protein pair) * L(longest seq) * Emb dim\n",
    "        padded_pairs = batch[0].to(device)\n",
    "        labels = batch[1].to(device)\n",
    "        \n",
    "        #split data into protA and protB\n",
    "        gosetA_batch = padded_pairs[:,0]\n",
    "        gosetB_batch = padded_pairs[:,1]\n",
    "        \n",
    "        vector_protA = torch.sum(gosetA_batch, dim = 1)\n",
    "        vector_protB = torch.sum(gosetB_batch, dim = 1)\n",
    "        \n",
    "        cosine = nn.CosineSimilarity(dim=1, eps=1e-6)\n",
    "        predictions = cosine(vector_protA, vector_protB)\n",
    "        loss = criterion(predictions, labels)\n",
    "        acc = binary_accuracy(predictions, labels)\n",
    "            \n",
    "        epoch_loss += loss.item()\n",
    "        epoch_acc += acc.item()\n",
    "            \n",
    "        pred = pred + list(predictions.cpu().data.numpy())\n",
    "        lab = lab + list(labels.cpu().data.numpy())\n",
    "\n",
    "    return epoch_loss / len(iterator), epoch_acc / len(iterator), lab, pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afaf80ac-f0c1-4038-815f-e0de2d48a758",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 129/129 [00:31<00:00,  4.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01 | Epoch Time: 0m 35s \tTrain Loss: 0.599 | Train Acc: 67.70% \t Val. Loss: 0.560 |  Val. Acc: 71.40% \t Roc Train: 0.736 \t Roc Valid: 0.821 ,   0.0001 --LR\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 22%|█████████████████▌                                                               | 28/129 [00:06<00:29,  3.45it/s]"
     ]
    }
   ],
   "source": [
    "model_name = \"model.pt\"\n",
    "\n",
    "writer = SummaryWriter(flush_secs=14)\n",
    "N_EPOCHS = 30\n",
    "best_roc_val = float('-inf')\n",
    "\n",
    "for epoch in range(N_EPOCHS):\n",
    "    \n",
    "\n",
    "    start_time = time.time()\n",
    "    train_loss, train_acc, roc_train = train(model, train_grt, optimizer, criterion, torch_vers = False)\n",
    "    valid_loss, valid_acc, roc_val, _, _ = evaluate(model, val_grt, criterion, torch_vers = False)   \n",
    "    end_time = time.time()\n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "    if best_roc_val < roc_val:\n",
    "        best_roc_val = roc_val\n",
    "        torch.save(model.state_dict(),  model_name)\n",
    "\n",
    "    print_status(epoch, epoch_mins, epoch_secs, train_loss,\\\n",
    "                 train_acc, valid_loss, valid_acc, roc_train, roc_val, optimizer)\n",
    "    write_scalars_tensorboard(writer, train_loss, valid_loss, train_acc, valid_acc, epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbd9c10a-3d3b-49dc-8c67-427338fd32ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "#WRiTING THE PERFORMANCE ON THE TEST SET #\n",
    "\n",
    "model = TransformerGO_Scratch(MODEL_SIZE, NR_HEADS, NR_LAYERS, SIZE_FF, DROPOUT)\n",
    "#model =  GO_Sum_NN(MODEL_SIZE, F_1, F_2, F_3, F_4, DROPOUT)\n",
    "\n",
    "model.load_state_dict(torch.load(model_name))\n",
    "model = model.to(device)\n",
    "\n",
    "with open(\"training-results.txt\", \"a\") as myfile:\n",
    "    myfile.write(f\"\\n ### {model_name} ### \\n\")\n",
    "    myfile.write(f\"Train set: {len(train_set)}, Valid set: {len(valid_set)}, Test set: {len(test_set)} \\n\")\n",
    "    \n",
    "    valid_loss, valid_acc, roc_val, lab, pred = evaluate(model, test_grt, criterion, torch_vers = False)\n",
    "    myfile.write(f\" \\n valid_loss: {valid_loss}, valid_acc: {valid_acc}, roc_val: {roc_val} \\n\")\n",
    "    \n",
    "    valid_loss, valid_acc, lab, pred = evaluate_cosine(test_grt, criterion, torch_vers = False)\n",
    "    roc_val = metrics.roc_auc_score(lab, pred)\n",
    "    myfile.write(f\" \\n ***COSINE*** \\n valid_loss: {valid_loss}, valid_acc: {valid_acc}, roc_val: {roc_val} \\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sacred-finder",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fold nr:  0\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 129/129 [00:45<00:00,  2.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01 | Epoch Time: 0m 50s \tTrain Loss: 0.597 | Train Acc: 68.14% \t Val. Loss: 0.557 |  Val. Acc: 70.77% \t Roc Train: 0.742 \t Roc Valid: 0.795 ,   0.0001 --LR\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 29%|███████████████████████▊                                                         | 38/129 [00:13<00:36,  2.47it/s]"
     ]
    }
   ],
   "source": [
    "#CROSS FOLD VALIDATION EXPERIMENT#\n",
    "\n",
    "C_FOLD = 5\n",
    "N_EPOCHS = 20\n",
    "\n",
    "MODEL_SIZE = EMB_DIM\n",
    "NR_HEADS = 8\n",
    "NR_LAYERS = 3\n",
    "DROPOUT = 0.2\n",
    "SIZE_FF = 4 * MODEL_SIZE\n",
    "\n",
    "LR = 0.0001\n",
    "\n",
    "sz = len(all_dataset)\n",
    "fold_size = int(sz/C_FOLD)\n",
    "l = 0\n",
    "r = fold_size\n",
    "indexes = np.arange(sz)\n",
    "\n",
    "val_accs = []\n",
    "val_rocs = []\n",
    "wrong_eval = []\n",
    "wrong_eval_labels = []\n",
    "for i in range(0, C_FOLD):\n",
    "    print(\"Fold nr: \", i, end='\\r')\n",
    "    \n",
    "    model = TransformerGO_Scratch(MODEL_SIZE, NR_HEADS, NR_LAYERS, SIZE_FF, DROPOUT).to(device)\n",
    "    pytorch_total_params = sum(p.numel() for p in model.parameters())\n",
    "    optimizer = optim.Adam(model.parameters(), lr=LR) \n",
    "    criterion = nn.BCEWithLogitsLoss().to(device)\n",
    "    writer = SummaryWriter(flush_secs=14)\n",
    "    \n",
    "    val_subset = data.Subset(all_dataset, indexes[l:r])\n",
    "    c_val_grt = data.DataLoader(val_subset, **params, shuffle = False)\n",
    "    \n",
    "    train_subset = data.Subset(all_dataset, np.concatenate([indexes[0:l], indexes[r:sz]]))\n",
    "    c_train_grt = data.DataLoader(train_subset, **params, shuffle = True)\n",
    "    \n",
    "    l += fold_size\n",
    "    r += fold_size\n",
    "\n",
    "    best_valid_roc = float('-inf')\n",
    "    best_valid_acc = float('-inf')\n",
    "    temp_w_eval = []\n",
    "    for epoch in range(N_EPOCHS):\n",
    "\n",
    "        start_time = time.time()\n",
    "        train_loss, train_acc, roc_train = train(model, c_train_grt, optimizer, criterion)\n",
    "        valid_loss, valid_acc, roc_val, _, _ = evaluate(model, c_val_grt, criterion)\n",
    "        end_time = time.time()\n",
    "        epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "        print_status(epoch, epoch_mins, epoch_secs, train_loss,\\\n",
    "                 train_acc, valid_loss, valid_acc, roc_train, roc_val, optimizer)\n",
    "        write_scalars_tensorboard(writer, train_loss, valid_loss, train_acc, valid_acc, epoch)\n",
    "        \n",
    "        best_valid_roc = max(best_valid_roc, roc_val)\n",
    "        best_valid_acc = max(best_valid_acc, valid_acc)\n",
    "\n",
    "    val_rocs.append(best_valid_roc)\n",
    "    val_accs.append(best_valid_acc)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "downtown-amplifier",
   "metadata": {},
   "outputs": [],
   "source": [
    "#saving the results from the 5 cross validation experiment\n",
    "with open(\"5cv_roc_\" + neg_path[-24:][:4] + neg_path[-5:] + '.pkl', \"wb\") as fp:\n",
    "    pickle.dump(val_rocs, fp)\n",
    "with open(\"5cv_acc_\" + neg_path[-24:][:4] + neg_path[-5:] + '.pkl', \"wb\") as fp:\n",
    "    pickle.dump(val_accs, fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2addcb2e-79c8-4e53-bc81-0e615698f187",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:ppi-phd]",
   "language": "python",
   "name": "conda-env-ppi-phd-py"
  },
  "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}