notebooks/finetuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e30ea18e-6b5a-47d4-b7a4-1330804b5602",
   "metadata": {},
   "source": [
    "# Finetune on experimental data\n",
    "This notebook demonstrates how to finetune METL models on experimental data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "432eebaf-00b8-42bf-b927-fd651e6ab94d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-16T22:51:36.573559Z",
     "start_time": "2024-02-16T22:51:36.569490Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c566507e-1012-4415-82ba-7498950e0b6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "# define the name of the project root directory\n",
    "project_root_dir_name = \"metl\"\n",
    "\n",
    "# find the project root by checking each parent directory\n",
    "current_dir = os.getcwd()\n",
    "while os.path.basename(current_dir) != project_root_dir_name and current_dir != os.path.dirname(current_dir):\n",
    "    current_dir = os.path.dirname(current_dir)\n",
    "\n",
    "# change the current working directory to the project root directory\n",
    "if os.path.basename(current_dir) == project_root_dir_name:\n",
    "    os.chdir(current_dir)\n",
    "else:\n",
    "    print(\"project root directory not found\")\n",
    "    \n",
    "# add the project code folder to the system path so imports work\n",
    "module_path = os.path.abspath(\"code\")\n",
    "if module_path not in sys.path:\n",
    "    sys.path.append(module_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19876208-66f9-46b5-8f50-8e798fa815a4",
   "metadata": {},
   "source": [
    "# Acquire an experimental dataset\n",
    "\n",
    "For demonstration purposes, this repository contains the [avGFP dataset](../data/dms_data/avgfp) from [Sarkisyan et al. (2016)](https://doi.org/10.1038/nature17995). \n",
    "See the [metl-pub](https://github.com/gitter-lab/metl-pub) repository to access the other experimental datasets we used in our preprint.\n",
    "See the README in the [dms_data](../data/dms_data) directory for information about how to use your own experimental dataset. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6abf8b1-aa2d-4055-9184-d962ba0d4582",
   "metadata": {},
   "source": [
    "# Acquire a pretrained model\n",
    "Pretrained METL models are available in the [metl-pretrained](https://github.com/gitter-lab/metl-pretrained) repository. You can use one of those, or you can pretrain your own METL model (see [pretraining.ipynb](pretraining.ipynb)). \n",
    "\n",
    "For demonstration purposes, we include a pretrained avGFP METL-Local model from the [metl-pretrained](https://github.com/gitter-lab/metl-pretrained) repository in the [pretrained_models](../pretrained_models) directory. This model is `METL-L-2M-3D-GFP` (UUID: `Hr4GNHws`).\n",
    "It is the avGFP METL-Local source model we used for the analysis in our preprint.\n",
    "\n",
    "We will show how to finetune this model using the [experimental avGFP dataset](../data/dms_data/avgfp)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23a30235-357a-4326-a4ff-77ab26eb5d7f",
   "metadata": {},
   "source": [
    "# Training arguments\n",
    "\n",
    "The script for finetuning on experimental data is [train_target_model.py](train_target_model.py). This script has a number of arguments you can view by uncommenting and running the below cell. There are additional arguments related to architecture that won't show up if you run the command, but you can view them in [models.py](../code/models.py) in the `TransferModel` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "bca8aeea-3dc3-47eb-915c-d80132be8fef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !python code/train_target_model.py -h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ec8c31b-2da2-4ba7-9f4e-39e30dce8056",
   "metadata": {},
   "source": [
    "We set up finetuning arguments for this example in [finetune_avgfp_local.txt](../args/pretrain_avgfp_local.txt) in the [args](../args) directory. This argument file can be used directly with [train_target_model.py](train_target_model.py) by calling the command `!python code/train_target_model.py @args/finetune_avgfp_local.txt` (we do this in the next section).\n",
    "\n",
    "Uncomment and run the cell below to view the contents of the argument file. The sections below will walk through and explain the key arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a06a897f-877d-4e41-9bee-4d3eabeead7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# with open(\"args/finetune_avgfp_local.txt\", \"r\") as file:\n",
    "#     contents = file.read()\n",
    "#     print(contents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2610124-fa2c-4709-98fc-bae51b258338",
   "metadata": {},
   "source": [
    "## Dataset arguments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f56ee90-90be-41fa-bc99-c13f94e14976",
   "metadata": {},
   "source": [
    "\n",
    "Specify the dataset name and the train/val/test split. The dataset must be defined in [datasets.yml](../data/dms_data/datasets.yml). For demonstration purposes, we are using one of the reduced dataset size splits with a dataset size of 160 (train size of 128).\n",
    "```\n",
    "--ds_name\n",
    "avgfp\n",
    "--split_dir\n",
    "data/dms_data/avgfp/splits/resampled/resampled_ds160_val0.2_te0.1_w1abc2f4e9a64_s1_r8099/resampled_ds160_val0.2_te0.1_w1abc2f4e9a64_s1_r8099_rep_0\n",
    "```\n",
    "\n",
    "Specify the names of the train, validation, and test set files in the split directory. Using \"auto\" for the test_name will select the super test set (\"stest.txt\") if it exists in the split directory, otherwise it will use the standard test set (\"test.txt\").\n",
    "\n",
    "```\n",
    "--train_name\n",
    "train\n",
    "--val_name\n",
    "val\n",
    "--test_name\n",
    "test\n",
    "```\n",
    "\n",
    "The name of the target column in the dataset dataframe. The model will be finetuned to predict the score in this column.\n",
    "\n",
    "```\n",
    "--target_names\n",
    "score\n",
    "```\n",
    "\n",
    "The METL-Local model we are finetuning uses 3D structure-based relative position embeddings, so we need to specify the PDB filename. This PDB file is in the [data/pdb_files](../data/pdb_files) directory, which the script checks by default, so there is no need to specify the full path. You can also just specify \"auto\" to use the PDB file defined for this dataset in [datasets.yml](../data/dms_data/datasets.yml).\n",
    "\n",
    "```\n",
    "--pdb_fn\n",
    "1gfl_cm.pdb\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "890cea13-feae-4e54-bf0f-dcbe97f4409f",
   "metadata": {},
   "source": [
    "## Network architecture arguments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72ee9762-cae7-4e21-8435-f6dd49781b8c",
   "metadata": {},
   "source": [
    "For finetuning, we implemented a special model `transfer_model` that handles pretrained checkpoints with top nets. \n",
    "```\n",
    "--model_name\n",
    "transfer_model\n",
    "```\n",
    "\n",
    "The pretrained checkpoint can be a PyTorch checkpoint (.pt file) downloaded from the [metl-pretrained](https://github.com/gitter-lab/metl-pretrained) repository or a PyTorch Lightning checkpoint (.ckpt file) obtained from pretraining a model with this repository.\n",
    "```\n",
    "--pretrained_ckpt_path\n",
    "pretrained_models/Hr4GNHws.pt\n",
    "```\n",
    "\n",
    "The backbone cutoff determines where to cutoff the pretrained model and place the new prediction head. For METL-Local models, we recommend backbone cutoff -1, and for METL-Global models we recommend backbone cutoff -2. \n",
    "\n",
    "```\n",
    "--backbone_cutoff\n",
    "-1\n",
    "```\n",
    "\n",
    "The remaining arguments determine the encoding, which should be set to `int_seqs`, whether to use dropout after the backbone cutoff, and the architecture of the new top net. You can leave these values as-is to match what we did for the preprint.\n",
    "\n",
    "```\n",
    "--encoding\n",
    "int_seqs\n",
    "--dropout_after_backbone\n",
    "--dropout_after_backbone_rate\n",
    "0.5\n",
    "--top_net_type\n",
    "linear\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d94c112-9770-4a5f-93e0-acf4d9acae16",
   "metadata": {},
   "source": [
    "## Finetuning strategy arguments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bb96cb6-7815-4efa-9b6f-305df9bb3050",
   "metadata": {},
   "source": [
    "We implemented a dual-phase finetuning strategy. During the first phase, the backbone weights are frozen and only the top net is trained. During the second phase, all the network weights are unfrozen and trained at a reduced learning rate.\n",
    "\n",
    "The unfreeze_backbone_at_epoch argument determines the training epoch at which to unfreeze the backbone. We train the models for 500 epochs, so the backbone is unfrozen halfway through at epoch 250.\n",
    "\n",
    "```\n",
    "--finetuning\n",
    "--finetuning_strategy\n",
    "backbone\n",
    "--unfreeze_backbone_at_epoch\n",
    "250\n",
    "--backbone_always_align_lr\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f51d4584-a0ce-45c3-8fb7-8c34d3a984c3",
   "metadata": {},
   "source": [
    "## Optimization arguments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d90d10e8-21f3-4b9e-8134-99cb053bef13",
   "metadata": {},
   "source": [
    "Basic optimizer arguments include the batch size, learning rate, and maximum number of epochs to train for. Unless early stopping is enabled, the model will train for the given number of epochs. \n",
    "\n",
    "```\n",
    "--optimizer\n",
    "adamw\n",
    "--weight_decay\n",
    "0.1\n",
    "--batch_size\n",
    "128\n",
    "--learning_rate\n",
    "0.001\n",
    "--max_epochs\n",
    "500\n",
    "--gradient_clip_val\n",
    "0.5\n",
    "```\n",
    "\n",
    "The learning rate scheduler we used for finetuning is a dual phase learning rate schedule that matches the dual phase finetuning strategy. Each phase has a linear learning rate warmup for 1% of the total steps in that phase. There is also a cosine decay for the learning rate for each phase. The phase 2 learning rate is 10% of the phase 1 learning rate.\n",
    "\n",
    "```\n",
    "--lr_scheduler\n",
    "dual_phase_warmup_constant_cosine_decay\n",
    "--warmup_steps\n",
    ".01\n",
    "--phase2_lr_ratio\n",
    "0.1\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16327f53-7beb-412e-a925-12884e66d70b",
   "metadata": {},
   "source": [
    "## Logging arguments"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "132db93c-85e6-4658-a31e-9b103df34cb7",
   "metadata": {},
   "source": [
    "We have built in functionality for tracking model training with Weights & Biases. If you have a Weights and Biases account, set the environment variable `WANDB_API_KEY` to your API key and set the flag `--use_wandb` instead of `--no_use_wandb` below.\n",
    "\n",
    "```\n",
    "--no_use_wandb\n",
    "--wandb_project\n",
    "metl-target\n",
    "--wandb_online\n",
    "--experiment\n",
    "default\n",
    "```\n",
    "\n",
    "The below argument determines where to place the log directory locally.\n",
    "```\n",
    "--log_dir_base\n",
    "output/training_logs\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53a2fda3-6dfa-46d5-ad3d-3055eda0b29a",
   "metadata": {},
   "source": [
    "# Running training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d3d8d23-9d54-4888-842d-4fc8fd843b40",
   "metadata": {},
   "source": [
    "All the arguments described above are contained in [finetune_avgfp_local.txt](../args/pretrain_avgfp_local.txt), which can be fed directly into [train_target_model.py](train_target_model.py).\n",
    "\n",
    "PyTorch Lightning has a built-in progress bar that is convenient for seeing training progress, but it does not display correctly in Jupyter when calling the script with `!python`. We are going to disable the progress bar for by setting the flag `--enable_progress_bar false`. Instead, we implemented a simple print statement to track training progress, which we will enable with the flag `--enable_simple_progress_messages`. \n",
    "\n",
    "The [train_target_model.py](../code/train_target_model.py) script can support running on Apple Silicon with acceleration via MPS, but the version of PyTorch used in this environment is slightly outdated and does not support all MPS operations, so MPS support has been disabled. The script will run on GPU via CUDA if available, otherwise it will use CPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "977b4d8d-4662-4e03-955c-dc4a8ae7c1dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Random seed not specified, using: 345050601\n",
      "Global seed set to 345050601\n",
      "Created model UUID: inqx2jYi\n",
      "Created log directory: output/training_logs/inqx2jYi\n",
      "Final UUID: inqx2jYi\n",
      "Final log directory: output/training_logs/inqx2jYi\n",
      "Trainer already configured with model summary callbacks: [<class 'pytorch_lightning.callbacks.rich_model_summary.RichModelSummary'>]. Skipping setting a default `ModelSummary` callback.\n",
      "GPU available: True (mps), used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "Loading `train_dataloader` to estimate number of stepping batches.\n",
      "/opt/homebrew/Caskroom/mambaforge/base/envs/metl/lib/python3.9/site-packages/pytorch_lightning/trainer/trainer.py:1892: PossibleUserWarning: The number of training batches (1) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.\n",
      "  rank_zero_warn(\n",
      "Number of training steps is 50\n",
      "Number of warmup steps is 0.5\n",
      "Second warmup phase starts at step 25\n",
      "total_steps 50\n",
      "phase1_total_steps 25\n",
      "phase2_total_steps 25\n",
      "┏━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━┳━━━━━━━━━━━━┓\n",
      "┃\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mName                  \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mType        \u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35mParams\u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35m  In sizes\u001b[0m\u001b[1;35m \u001b[0m┃\u001b[1;35m \u001b[0m\u001b[1;35m Out sizes\u001b[0m\u001b[1;35m \u001b[0m┃\n",
      "┡━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━╇━━━━━━━━━━━━┩\n",
      "│\u001b[2m \u001b[0m\u001b[2m0\u001b[0m\u001b[2m \u001b[0m│ model                  │ TransferMod… │  2.4 M │\u001b[37m \u001b[0m\u001b[37m[128, 237]\u001b[0m\u001b[37m \u001b[0m│\u001b[37m \u001b[0m\u001b[37m  [128, 1]\u001b[0m\u001b[37m \u001b[0m│\n",
      "│\u001b[2m \u001b[0m\u001b[2m1\u001b[0m\u001b[2m \u001b[0m│ model.model            │ SequentialW… │  2.4 M │\u001b[37m \u001b[0m\u001b[37m[128, 237]\u001b[0m\u001b[37m \u001b[0m│\u001b[37m \u001b[0m\u001b[37m  [128, 1]\u001b[0m\u001b[37m \u001b[0m│\n",
      "│\u001b[2m \u001b[0m\u001b[2m2\u001b[0m\u001b[2m \u001b[0m│ model.model.backbone   │ SequentialW… │  2.4 M │\u001b[37m \u001b[0m\u001b[37m[128, 237]\u001b[0m\u001b[37m \u001b[0m│\u001b[37m \u001b[0m\u001b[37m[128, 256]\u001b[0m\u001b[37m \u001b[0m│\n",
      "│\u001b[2m \u001b[0m\u001b[2m3\u001b[0m\u001b[2m \u001b[0m│ model.model.dropout    │ Dropout      │      0 │\u001b[37m \u001b[0m\u001b[37m[128, 256]\u001b[0m\u001b[37m \u001b[0m│\u001b[37m \u001b[0m\u001b[37m[128, 256]\u001b[0m\u001b[37m \u001b[0m│\n",
      "│\u001b[2m \u001b[0m\u001b[2m4\u001b[0m\u001b[2m \u001b[0m│ model.model.flatten    │ Flatten      │      0 │\u001b[37m \u001b[0m\u001b[37m[128, 256]\u001b[0m\u001b[37m \u001b[0m│\u001b[37m \u001b[0m\u001b[37m[128, 256]\u001b[0m\u001b[37m \u001b[0m│\n",
      "│\u001b[2m \u001b[0m\u001b[2m5\u001b[0m\u001b[2m \u001b[0m│ model.model.prediction │ Linear       │    257 │\u001b[37m \u001b[0m\u001b[37m[128, 256]\u001b[0m\u001b[37m \u001b[0m│\u001b[37m \u001b[0m\u001b[37m  [128, 1]\u001b[0m\u001b[37m \u001b[0m│\n",
      "│\u001b[2m \u001b[0m\u001b[2m6\u001b[0m\u001b[2m \u001b[0m│ test_pearson           │ PearsonCorr… │      0 │\u001b[37m \u001b[0m\u001b[37m         ?\u001b[0m\u001b[37m \u001b[0m│\u001b[37m \u001b[0m\u001b[37m         ?\u001b[0m\u001b[37m \u001b[0m│\n",
      "│\u001b[2m \u001b[0m\u001b[2m7\u001b[0m\u001b[2m \u001b[0m│ test_spearman          │ SpearmanCor… │      0 │\u001b[37m \u001b[0m\u001b[37m         ?\u001b[0m\u001b[37m \u001b[0m│\u001b[37m \u001b[0m\u001b[37m         ?\u001b[0m\u001b[37m \u001b[0m│\n",
      "└───┴────────────────────────┴──────────────┴────────┴────────────┴────────────┘\n",
      "\u001b[1mTrainable params\u001b[0m: 257                                                           \n",
      "\u001b[1mNon-trainable params\u001b[0m: 2.4 M                                                     \n",
      "\u001b[1mTotal params\u001b[0m: 2.4 M                                                             \n",
      "\u001b[1mTotal estimated model params size (MB)\u001b[0m: 9                                       \n",
      "Starting sanity check...\n",
      "Sanity check complete.\n",
      "Starting training...\n",
      "Epoch     0: Train Loss =   2.460, Val Loss =   2.208\n",
      "Epoch     1: Train Loss =   2.387, Val Loss =   2.159\n",
      "Epoch     2: Train Loss =   2.434, Val Loss =   2.112\n",
      "Epoch     3: Train Loss =   2.257, Val Loss =   2.066\n",
      "Epoch     4: Train Loss =   2.278, Val Loss =   2.022\n",
      "Epoch     5: Train Loss =   2.185, Val Loss =   1.981\n",
      "Epoch     6: Train Loss =   2.236, Val Loss =   1.942\n",
      "Epoch     7: Train Loss =   2.142, Val Loss =   1.906\n",
      "Epoch     8: Train Loss =   2.106, Val Loss =   1.873\n",
      "Epoch     9: Train Loss =   2.078, Val Loss =   1.842\n",
      "Epoch    10: Train Loss =   2.020, Val Loss =   1.815\n",
      "Epoch    11: Train Loss =   2.038, Val Loss =   1.790\n",
      "Epoch    12: Train Loss =   1.971, Val Loss =   1.768\n",
      "Epoch    13: Train Loss =   1.989, Val Loss =   1.749\n",
      "Epoch    14: Train Loss =   1.924, Val Loss =   1.732\n",
      "Epoch    15: Train Loss =   1.921, Val Loss =   1.718\n",
      "Epoch    16: Train Loss =   1.874, Val Loss =   1.707\n",
      "Epoch    17: Train Loss =   1.905, Val Loss =   1.698\n",
      "Epoch    18: Train Loss =   1.871, Val Loss =   1.691\n",
      "Epoch    19: Train Loss =   1.830, Val Loss =   1.685\n",
      "Epoch    20: Train Loss =   1.910, Val Loss =   1.682\n",
      "Epoch    21: Train Loss =   1.858, Val Loss =   1.679\n",
      "Epoch    22: Train Loss =   1.869, Val Loss =   1.678\n",
      "Epoch    23: Train Loss =   1.868, Val Loss =   1.677\n",
      "Epoch    24: Train Loss =   1.858, Val Loss =   1.677\n",
      "Epoch    25: Train Loss =   1.813, Val Loss =   1.677\n",
      "Epoch    26: Train Loss =   1.857, Val Loss =   1.650\n",
      "Epoch    27: Train Loss =   1.850, Val Loss =   1.624\n",
      "Epoch    28: Train Loss =   1.813, Val Loss =   1.597\n",
      "Epoch    29: Train Loss =   1.773, Val Loss =   1.570\n",
      "Epoch    30: Train Loss =   1.712, Val Loss =   1.544\n",
      "Epoch    31: Train Loss =   1.696, Val Loss =   1.519\n",
      "Epoch    32: Train Loss =   1.690, Val Loss =   1.494\n",
      "Epoch    33: Train Loss =   1.638, Val Loss =   1.470\n",
      "Epoch    34: Train Loss =   1.629, Val Loss =   1.447\n",
      "Epoch    35: Train Loss =   1.630, Val Loss =   1.424\n",
      "Epoch    36: Train Loss =   1.560, Val Loss =   1.403\n",
      "Epoch    37: Train Loss =   1.524, Val Loss =   1.383\n",
      "Epoch    38: Train Loss =   1.478, Val Loss =   1.364\n",
      "Epoch    39: Train Loss =   1.476, Val Loss =   1.348\n",
      "Epoch    40: Train Loss =   1.396, Val Loss =   1.332\n",
      "Epoch    41: Train Loss =   1.428, Val Loss =   1.320\n",
      "Epoch    42: Train Loss =   1.355, Val Loss =   1.309\n",
      "Epoch    43: Train Loss =   1.369, Val Loss =   1.300\n",
      "Epoch    44: Train Loss =   1.394, Val Loss =   1.293\n",
      "Epoch    45: Train Loss =   1.346, Val Loss =   1.288\n",
      "Epoch    46: Train Loss =   1.340, Val Loss =   1.285\n",
      "Epoch    47: Train Loss =   1.335, Val Loss =   1.283\n",
      "Epoch    48: Train Loss =   1.356, Val Loss =   1.282\n",
      "Epoch    49: Train Loss =   1.246, Val Loss =   1.282\n",
      "`Trainer.fit` stopped: `max_epochs=50` reached.\n",
      "Restoring states from the checkpoint path at output/training_logs/inqx2jYi/checkpoints/epoch=49-step=50.ckpt\n",
      "Loaded model weights from checkpoint at output/training_logs/inqx2jYi/checkpoints/epoch=49-step=50.ckpt\n",
      "Starting testing...\n",
      "Testing complete.\n",
      "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
      "┃\u001b[1m \u001b[0m\u001b[1m       Test metric       \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m      DataLoader 0       \u001b[0m\u001b[1m \u001b[0m┃\n",
      "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
      "│\u001b[36m \u001b[0m\u001b[36m        test_loss        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m    1.324419379234314    \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test_pearson       \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.6309216618537903    \u001b[0m\u001b[35m \u001b[0m│\n",
      "│\u001b[36m \u001b[0m\u001b[36m      test_spearman      \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m   0.6115598678588867    \u001b[0m\u001b[35m \u001b[0m│\n",
      "└───────────────────────────┴───────────────────────────┘\n",
      "Restoring states from the checkpoint path at output/training_logs/inqx2jYi/checkpoints/epoch=49-step=50.ckpt\n",
      "Loaded model weights from checkpoint at output/training_logs/inqx2jYi/checkpoints/epoch=49-step=50.ckpt\n",
      "Starting prediction...\n",
      "Prediction complete.\n",
      "saving a scatter plot for set: train (128 variants)\n",
      "saving a scatter plot for set: val (32 variants)\n",
      "saving a scatter plot for set: test (4655 variants)\n",
      "            mse  pearsonr        r2  spearmanr\n",
      "set                                           \n",
      "train  1.416450  0.704978 -0.214095   0.669300\n",
      "val    1.281733  0.684139 -0.188552   0.650536\n",
      "test   1.324419  0.630922 -0.189759   0.611560\n"
     ]
    }
   ],
   "source": [
    "!python code/train_target_model.py @args/finetune_avgfp_local.txt --enable_progress_bar false --enable_simple_progress_messages --max_epochs 50 --unfreeze_backbone_at_epoch 25"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f33fc407-6ab1-45e3-8e6a-9b717dca7f00",
   "metadata": {},
   "source": [
    "# Additional recommendations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39c8e0e5-8bb5-4200-ab45-e559b0f20896",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "## Model selection\n",
    "\n",
    "Selecting the model from the epoch with the lowest validation set loss can help prevent overfitting. It requires having a big enough validation set that provides an accurate estimate of performance. \n",
    "\n",
    "We enabled model selection if the validation set size was ≥ 32 for METL-Local and ≥ 128 for METL-Global. We found the optimization was more stable for METL-Local than METL-Global, thus smaller validation sets were still reliable. \n",
    "\n",
    "Enable model selection by setting argument `--ckpt_monitor val`.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b18f773b-8209-4993-b3f0-994b0ab2b133",
   "metadata": {},
   "source": [
    "## Backbone cutoff for METL-Global\n",
    "Finetuning METL-Global is largely the same as METL-Local, except we recommend using a different threshold for model selection (see above), as well as a different backbone cutoff.\n",
    "\n",
    "For METL-Local, we set `--backbone_cutoff -1`, which attaches the new prediction head immediately after the final fully connected layer. \n",
    "\n",
    "For METL-Global, we recommend setting `--backbone_cutoff -2`, which attaches the new prediction head immediately after the global pooling layer. We found this resulted in better finetuning performance for METL-Global."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a591eb8-3d5e-437f-9189-3c0834f7f447",
   "metadata": {},
   "source": [
    "# Running inference using finetuned model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af85ff8f-1a30-4ba2-bf3b-967a773e0e80",
   "metadata": {},
   "source": [
    "The PyTorch Lightning framework supports inference, but while we put together a working example, we recommend converting the PyTorch Lightning checkpoint to pure PyTorch and using the [metl-pretrained](https://github.com/gitter-lab/metl-pretrained) package to run inference in pure PyTorch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1acca5d1-1bca-4c3f-b9d3-56525cf11186",
   "metadata": {},
   "source": [
    "## Convert to PyTorch\n",
    "Lightning checkpoints are compatible with pure pytorch, but they may contain additional items that are not needed for inference. This script loads the checkpoint and saves a smaller checkpoint with just the model weights and hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "63d8ce0a-5534-406f-90b6-6c155cb6ea9c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing checkpoint: output/training_logs/inqx2jYi/checkpoints/epoch=49-step=50.ckpt\n",
      "Saving converted checkpoint to: output/training_logs/inqx2jYi/checkpoints/inqx2jYi.pt\n"
     ]
    }
   ],
   "source": [
    "# the Lightning checkpoint from the finetuning we performed above\n",
    "ckpt_fn = \"output/training_logs/inqx2jYi/checkpoints/epoch=49-step=50.ckpt\"\n",
    "\n",
    "# run the conversion script\n",
    "!python code/convert_ckpt.py --ckpt_path $ckpt_fn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b562aa-663a-4b0d-a719-e85555cf875d",
   "metadata": {},
   "source": [
    "## Load checkpoint with metl-pretrained package\n",
    "Install the `metl` package according to the instructions in the [metl-pretrained](https://github.com/gitter-lab/metl-pretrained) repository. You can install it directly into this environment. Load the model using `metl.get_from_checkpoint()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ff35ce69-97ed-4a5a-b082-f197aae1addc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initialized PDB bucket matrices in: 0.000\n",
      "Initialized PDB bucket matrices in: 0.000\n"
     ]
    }
   ],
   "source": [
    "import metl\n",
    "import torch\n",
    "import utils\n",
    "\n",
    "checkpoint_path = \"output/training_logs/inqx2jYi/checkpoints/inqx2jYi.pt\"\n",
    "model, data_encoder = metl.get_from_checkpoint(checkpoint_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7efa0346-d84b-42ca-ab4b-ed4ae3120933",
   "metadata": {},
   "source": [
    "## Run inference with pure PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "33202a3e-5e35-4b55-8738-3b58ad04796e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.2634],\n",
      "        [-0.3954],\n",
      "        [-0.6807]])\n"
     ]
    }
   ],
   "source": [
    "# load the GFP wild-type sequence and the PDB file (needed for 3D RPE)\n",
    "datasets = utils.load_dataset_metadata()\n",
    "wt = datasets[\"avgfp\"][\"wt_aa\"]\n",
    "pdb_fn = datasets[\"avgfp\"][\"pdb_fn\"]\n",
    "\n",
    "# some example GFP variants to compute the scores for\n",
    "variants = [\"E3K,G102S\",\n",
    "            \"T36P,S203T,K207R\",\n",
    "            \"V10A,D19G,F25S,E113V\"]\n",
    "\n",
    "encoded_variants = data_encoder.encode_variants(wt, variants)\n",
    "\n",
    "# set model to eval mode\n",
    "model.eval()\n",
    "\n",
    "# no need to compute gradients for inference\n",
    "with torch.no_grad():\n",
    "    # note we are specifying the pdb_fn because this model uses 3D RPE\n",
    "    predictions = model(torch.tensor(encoded_variants), pdb_fn=pdb_fn)\n",
    "\n",
    "print(predictions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed09dc55-6c95-4509-82aa-448eded0e166",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}