04.train_final_ensemble.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e10590fd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: TORCH_HOME=/ayb/vol2/home/dumerenkov/torch_hub\n"
     ]
    }
   ],
   "source": [
    "%env TORCH_HOME=/ayb/vol2/home/dumerenkov/torch_hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "14e770d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import Bio\n",
    "from Bio import SeqIO\n",
    "import os\n",
    "import torch\n",
    "import math\n",
    "import esm\n",
    "\n",
    "from torch import nn\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from sklearn.model_selection import train_test_split\n",
    "import scipy\n",
    "from scipy import stats\n",
    "\n",
    "import gc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6de1bfa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "81437c93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "\n",
    "HIDDEN_UNITS_POS_CONTACT = 5\n",
    "class ESMForSingleMutationPosConcat(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.esm2, _ = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self.fc1 = nn.Linear(1280 * 2, HIDDEN_UNITS_POS_CONTACT)\n",
    "        self.fc2 = nn.Linear(HIDDEN_UNITS_POS_CONTACT, 1)\n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):\n",
    "        outputs1 = self.esm2.forward(token_ids1, repr_layers=[33])[\n",
    "            'representations'][33]\n",
    "        outputs2 = self.esm2.forward(token_ids2, repr_layers=[33])[\n",
    "            'representations'][33]\n",
    "        outputs1_pos = outputs1[:, pos + 1]\n",
    "        outputs2_pos = outputs2[:, pos + 1]\n",
    "        outputs_pos_concat = torch.cat((outputs1_pos, outputs2_pos), 2)\n",
    "        fc1_outputs = F.relu(self.fc1(outputs_pos_concat))\n",
    "        logits = self.fc2(fc1_outputs)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8f39f27b",
   "metadata": {},
   "outputs": [],
   "source": [
    "HIDDEN_UNITS_POS_OUTER = 5\n",
    "class ESMForSingleMutationPosOuter(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.esm2, _ = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self._freeze_esm2_layers()\n",
    "        self.fc1 = nn.Linear(1280 * 1280, HIDDEN_UNITS_POS_OUTER)\n",
    "        self.fc2 = nn.Linear(HIDDEN_UNITS_POS_OUTER, 1)\n",
    "\n",
    "    def _freeze_esm2_layers(self):\n",
    "        total_blocks = 33\n",
    "        initial_layers = 2\n",
    "        layers_per_block = 16\n",
    "        num_freeze_blocks = total_blocks - 3\n",
    "        for _, param in list(self.esm2.named_parameters())[\n",
    "            :initial_layers + layers_per_block * num_freeze_blocks]:\n",
    "            param.requires_grad = False\n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):\n",
    "        outputs1 = self.esm2.forward(token_ids1, repr_layers=[33])[\n",
    "            'representations'][33]\n",
    "        outputs2 = self.esm2.forward(token_ids2, repr_layers=[33])[\n",
    "            'representations'][33]\n",
    "        outputs1_pos = outputs1[:, pos + 1]\n",
    "        outputs2_pos = outputs2[:, pos + 1]\n",
    "        outer_prod = outputs1_pos.unsqueeze(3) @ outputs2_pos.unsqueeze(2)\n",
    "        outer_prod_view = outer_prod.view(outer_prod.shape[0], outer_prod.shape[1], -1)\n",
    "        fc1_outputs = F.relu(self.fc1(outer_prod_view))\n",
    "        logits = self.fc2(fc1_outputs)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9aede66e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ESMForSingleMutation_pos(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__() \n",
    "        self.esm1v, self.esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()        \n",
    "        self.classifier = nn.Linear(1280, 1)\n",
    "        self.const1 = torch.nn.Parameter(torch.ones((1,1280)))\n",
    "        self.const2 = torch.nn.Parameter(-1 * torch.ones((1,1280)))\n",
    "        \n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):                \n",
    "        outputs1 = self.esm1v.forward(token_ids1, repr_layers=[33])['representations'][33]\n",
    "        outputs2 = self.esm1v.forward(token_ids2, repr_layers=[33])['representations'][33]\n",
    "        outputs = self.const1 * outputs1[:,pos + 1,:] + self.const2 * outputs2[:,pos + 1,:]        \n",
    "        logits = self.classifier(outputs)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "838b0855",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ESMForSingleMutation_cls(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__() \n",
    "        self.esm1v, self.esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()        \n",
    "        self.classifier = nn.Linear(1280, 1)\n",
    "        self.const1 = torch.nn.Parameter(torch.ones((1,1280)))\n",
    "        self.const2 = torch.nn.Parameter(-1 * torch.ones((1,1280)))\n",
    "        \n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):                \n",
    "        outputs1 = self.esm1v.forward(token_ids1, repr_layers=[33])['representations'][33]\n",
    "        outputs2 = self.esm1v.forward(token_ids2, repr_layers=[33])['representations'][33]\n",
    "        outputs = self.const1 * outputs1[:,0,:] + self.const2 * outputs2[:,0,:]        \n",
    "        logits = self.classifier(outputs.unsqueeze(0))\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1df241db",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ESMForSingleMutation_pos_cat_cls(nn.Module):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__() \n",
    "        self.esm1v, self.esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()        \n",
    "        self.classifier = nn.Linear(1280*2, 1)\n",
    "        self.const1 = torch.nn.Parameter(torch.ones((1,1280)))\n",
    "        self.const2 = torch.nn.Parameter(-1 * torch.ones((1,1280)))\n",
    "        \n",
    "\n",
    "    def forward(self, token_ids1, token_ids2, pos):                \n",
    "        outputs1 = self.esm1v.forward(token_ids1, repr_layers=[33])['representations'][33]\n",
    "        outputs2 = self.esm1v.forward(token_ids2, repr_layers=[33])['representations'][33]\n",
    "        cls_out = self.const1 * outputs1[:,0,:] + self.const2 * outputs2[:,0,:]\n",
    "        pos_out = self.const1 * outputs1[:,pos+1,:] + self.const2 * outputs2[:,pos+1,:]\n",
    "        outputs = torch.cat([cls_out.unsqueeze(0), pos_out], axis = -1)        \n",
    "        logits = self.classifier(outputs)\n",
    "        return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1a1cb550",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProteinDataset(Dataset):\n",
    "    def __init__(self, df):\n",
    "        self.df = df\n",
    "        _, esm1v_alphabet = esm.pretrained.esm2_t33_650M_UR50D()\n",
    "        self.esm1v_batch_converter = esm1v_alphabet.get_batch_converter()\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        _, _, esm1b_batch_tokens1 = self.esm1v_batch_converter([('' , ''.join(self.df.iloc[idx]['wt_seq'])[:1022])])\n",
    "        _, _, esm1b_batch_tokens2 = self.esm1v_batch_converter([('' , ''.join(self.df.iloc[idx]['mut_seq'])[:1022])])\n",
    "        pos = self.df.iloc[idx]['pos']\n",
    "        return esm1b_batch_tokens1, esm1b_batch_tokens2, pos, torch.unsqueeze(torch.FloatTensor([self.df.iloc[idx]['ddg']]), 0)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "642d008e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    tr_loss, tr_accuracy = 0, 0\n",
    "    nb_tr_examples, nb_tr_steps = 0, 0\n",
    "    tr_preds, tr_labels = [], []\n",
    "    model.train()\n",
    "    \n",
    "    for idx, batch in enumerate(training_loader):\n",
    "        input_ids1, input_ids2, pos, labels = batch            \n",
    "        input_ids1 = input_ids1[0].to(device)\n",
    "        input_ids2 = input_ids2[0].to(device)\n",
    "        logits = model(token_ids1 = input_ids1, token_ids2 = input_ids2, pos = pos).to('cpu')        \n",
    "        loss = torch.nn.functional.mse_loss(logits, labels)\n",
    "        tr_loss += loss.item()\n",
    "\n",
    "        nb_tr_steps += 1\n",
    "        nb_tr_examples += labels.size(0)\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(\n",
    "            parameters=model.parameters(), max_norm=0.1\n",
    "        )\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "    epoch_loss = tr_loss / nb_tr_steps\n",
    "    print(f\"Training loss epoch: {epoch_loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f9dfa131",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training model ESMForSingleMutationPosOuter\n",
      "Training loss epoch: 3.876894674672059\n",
      "Training loss epoch: 3.1980062041550696\n",
      "Training loss epoch: 2.915139241010014\n",
      "Training model ESMForSingleMutationPosConcat\n",
      "Training loss epoch: 3.860952251269573\n",
      "Training loss epoch: 1.469412682013314\n",
      "Training loss epoch: 0.661039534414614\n",
      "Training model ESMForSingleMutation_pos_cat_cls\n",
      "Training loss epoch: 2.6397629074271167\n",
      "Training loss epoch: 0.8607995854492856\n",
      "Training loss epoch: 0.3215935788703277\n",
      "Training model ESMForSingleMutation_pos\n",
      "Training loss epoch: 2.5457445398761536\n",
      "Training loss epoch: 0.8276246989439202\n",
      "Training loss epoch: 0.2824712599465907\n",
      "Training model ESMForSingleMutation_cls\n",
      "Training loss epoch: 3.6326927324780582\n",
      "Training loss epoch: 2.246442528126786\n",
      "Training loss epoch: 0.9332105714671484\n"
     ]
    }
   ],
   "source": [
    "lr = 1e-5\n",
    "EPOCHS = 3\n",
    "device = 'cuda:1'\n",
    "\n",
    "models = ['ESMForSingleMutationPosOuter',\n",
    "          'ESMForSingleMutationPosConcat',\n",
    "          'ESMForSingleMutation_pos_cat_cls',  \n",
    "              'ESMForSingleMutation_pos', \n",
    "              'ESMForSingleMutation_cls']\n",
    "\n",
    "full_df = pd.read_csv('DATASETS/new_ds_with_folds.csv')\n",
    "\n",
    "preds = {n:[] for n in models} \n",
    "true = [None]*5\n",
    "\n",
    "for model_name in models:\n",
    "    model_class = globals()[model_name]\n",
    "    print(f'Training model {model_name}')\n",
    "    train_df = full_df\n",
    "    train_ds = ProteinDataset(train_df)\n",
    "        \n",
    "    model = model_class()                        \n",
    "    model.to(device) \n",
    "    optimizer = torch.optim.Adam(params=model.parameters(), lr=lr)\n",
    "    training_loader = DataLoader(train_ds, batch_size=1, num_workers = 2, shuffle = True)\n",
    "    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=lr, steps_per_epoch=len(training_loader), epochs=EPOCHS)\n",
    "        \n",
    "    for epoch in range(EPOCHS):\n",
    "        train(epoch)\n",
    "         \n",
    "    model.to('cpu')\n",
    "    \n",
    "    torch.save(model, 'weights/' + model_name)\n",
    "    \n",
    "    del model"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}