{ "cells": [ { "cell_type": "markdown", "id": "79031066-2ad7-43eb-8143-66180d8a83d4", "metadata": {}, "source": [ "### We aim to train reward models using the 5'UTR polysome profiling data provided by Sample et al. Human MPTA Dataset." ] }, { "cell_type": "markdown", "id": "c71a1e10-55ef-4d44-bdb0-09eb5f4e49ba", "metadata": {}, "source": [ "- Paper link: https://www.nature.com/articles/s41587-019-0164-5\n", "- Data link: https://www.ncbi.nlm.nih.gov/geo/query/acc.cgi?acc=GSE114002" ] }, { "cell_type": "markdown", "id": "8134eac4-d3ce-40a0-89e8-e5af45fab5e7", "metadata": {}, "source": [ "## 1- Dataset preprocessing" ] }, { "cell_type": "code", "execution_count": 1, "id": "66b976a1-fa4e-4a69-9d0f-4e1e1033c6b8", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:31:15.678510Z", "iopub.status.busy": "2023-11-16T20:31:15.678377Z", "iopub.status.idle": "2023-11-16T20:31:20.255734Z", "shell.execute_reply": "2023-11-16T20:31:20.254973Z", "shell.execute_reply.started": "2023-11-16T20:31:15.678493Z" }, "tags": [] }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n", "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33msarosavo\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n" ] }, { "data": { "text/plain": [ "True" ] }, "execution_count": 1, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import pandas as pd\n", "import numpy as np\n", "\n", "import wandb\n", "wandb.login(host=\"https://api.wandb.ai\")" ] }, { "cell_type": "markdown", "id": "3f2db34c", "metadata": {}, "source": [ "### Load the dataset" ] }, { "cell_type": "code", "execution_count": 2, "id": "75b48cd2-2d91-4276-82ad-bafe5896cea0", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:32:01.845969Z", "iopub.status.busy": "2023-11-16T20:32:01.845833Z", "iopub.status.idle": "2023-11-16T20:32:01.865882Z", "shell.execute_reply": "2023-11-16T20:32:01.865467Z", "shell.execute_reply.started": "2023-11-16T20:32:01.845956Z" } }, "outputs": [ { "data": { "text/html": [ "
\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
utrrl
0CCACTCGATTAACATGTTAACAACATACTCGTCCGGCCGATCAGCG...-1.949703
1CAAATCATGTGCAGCCCTGGCGACCGTACTGCGGTACAAGAAAGTA...-1.460671
2GTTATACTAGAAGAAACTTGAGATTATGGAGCAGTCCGTCAAGGAC...-1.781243
3CTTAGACAAAAACAACGCGCTTTCCAGTATGCGGAGCCTTGACGGT...-1.643680
4GTATCAAATCACGGCCAACCCGACGGAGTACCCCGCGTCGATGGTC...-1.061173
\n", "
" ], "text/plain": [ " utr rl\n", "0 CCACTCGATTAACATGTTAACAACATACTCGTCCGGCCGATCAGCG... -1.949703\n", "1 CAAATCATGTGCAGCCCTGGCGACCGTACTGCGGTACAAGAAAGTA... -1.460671\n", "2 GTTATACTAGAAGAAACTTGAGATTATGGAGCAGTCCGTCAAGGAC... -1.781243\n", "3 CTTAGACAAAAACAACGCGCTTTCCAGTATGCGGAGCCTTGACGGT... -1.643680\n", "4 GTATCAAATCACGGCCAACCCGACGGAGTACCCCGCGTCGATGGTC... -1.061173" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "df = pd.read_csv(\"artifacts/UTR-dataset:v0/dataset.csv.gz\")\n", "df = df.drop('Unnamed: 0', axis = 1)\n", "df.head()" ] }, { "cell_type": "markdown", "id": "94e3f181-0b5d-4d90-a51e-99586d961d20", "metadata": {}, "source": [ "## 2- Model training" ] }, { "cell_type": "code", "execution_count": 3, "id": "4075f692-e30b-4c2a-9726-68d5920dfb71", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:32:37.959007Z", "iopub.status.busy": "2023-11-16T20:32:37.958577Z", "iopub.status.idle": "2023-11-16T20:32:37.961963Z", "shell.execute_reply": "2023-11-16T20:32:37.961493Z", "shell.execute_reply.started": "2023-11-16T20:32:37.958989Z" }, "tags": [] }, "outputs": [], "source": [ "import wandb\n", "import torch\n", "import grelu\n", "import pandas as pd\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "81b752d3-65db-47e9-a776-964eb5eebb18", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:32:38.196732Z", "iopub.status.busy": "2023-11-16T20:32:38.196439Z", "iopub.status.idle": "2023-11-16T20:33:06.025636Z", "shell.execute_reply": "2023-11-16T20:33:06.025192Z", "shell.execute_reply.started": "2023-11-16T20:32:38.196717Z" }, "tags": [] }, "outputs": [ { "data": { "text/html": [ "Tracking run with wandb version 0.17.4" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "Run data is saved locally in /raid/home/zhaoy183/projects/RLfinetuning_Diffusion_Bioseq/tutorials/UTR/wandb/run-20240718_142535-zy01qfcw" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "Syncing run train to Weights & Biases (docs)
" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ " View project at https://wandb.ai/sarosavo/UTR-design" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ " View run at https://wandb.ai/sarosavo/UTR-design/runs/zy01qfcw" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/plain": [ "'train'" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "run = wandb.init(entity=None, project='UTR-design', job_type='training', name='train')\n", "run.name" ] }, { "cell_type": "markdown", "id": "947e97a7-054d-4c48-b31d-17b33929a62e", "metadata": {}, "source": [ "### Train/Val/Test split " ] }, { "cell_type": "code", "execution_count": 5, "id": "857c5cfd-cde7-49b0-9165-ce2bab44f484", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:33:07.617564Z", "iopub.status.busy": "2023-11-16T20:33:07.617399Z", "iopub.status.idle": "2023-11-16T20:33:07.636214Z", "shell.execute_reply": "2023-11-16T20:33:07.635719Z", "shell.execute_reply.started": "2023-11-16T20:33:07.617546Z" } }, "outputs": [], "source": [ "import grelu.data.preprocess\n", "from grelu.lightning import LightningModel\n", "from sklearn.model_selection import train_test_split" ] }, { "cell_type": "code", "execution_count": 6, "id": "c437e4bd-f6c6-4c81-bd20-82bd7513d9dc", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:33:07.637000Z", "iopub.status.busy": "2023-11-16T20:33:07.636847Z", "iopub.status.idle": "2023-11-16T20:33:07.667491Z", "shell.execute_reply": "2023-11-16T20:33:07.666956Z", "shell.execute_reply.started": "2023-11-16T20:33:07.636985Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "(260826, 2)\n", "(32603, 2)\n" ] } ], "source": [ "df_train, df_test = train_test_split(df, test_size=0.2)\n", "df_test, df_val = train_test_split(df_test, test_size=0.5)\n", "print(df_train.shape)\n", "print(df_test.shape)" ] }, { "cell_type": "markdown", "id": "d9be6cd6-8ad3-4199-a7ed-5bdc9115ea9a", "metadata": {}, "source": [ "### Start the training!" ] }, { "cell_type": "code", "execution_count": 7, "id": "03d03609-aaaf-41b1-ba3e-4dd6e5642f22", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:33:07.681604Z", "iopub.status.busy": "2023-11-16T20:33:07.681448Z", "iopub.status.idle": "2023-11-16T20:33:07.684849Z", "shell.execute_reply": "2023-11-16T20:33:07.684424Z", "shell.execute_reply.started": "2023-11-16T20:33:07.681590Z" }, "tags": [] }, "outputs": [], "source": [ "model_params = {\n", " 'model_type':'EnformerPretrainedModel',\n", " 'n_tasks': 1,\n", " 'n_transformers': 3,\n", "}\n", "\n", "train_params = {\n", " 'task':'regression',\n", " 'loss': 'MSE',\n", " 'lr':1e-4,\n", " 'logger': 'wandb',\n", " 'batch_size': 256,\n", " 'num_workers': 4,\n", " 'devices': [1],\n", " 'save_dir': 'experiment',\n", " 'optimizer': 'adam',\n", " 'max_epochs': 10,\n", " 'checkpoint': True,\n", "}" ] }, { "cell_type": "code", "execution_count": 8, "id": "7799e8da-d2d3-4ec5-8788-77c52b855fa4", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:33:07.688685Z", "iopub.status.busy": "2023-11-16T20:33:07.688541Z", "iopub.status.idle": "2023-11-16T20:33:11.503422Z", "shell.execute_reply": "2023-11-16T20:33:11.502819Z", "shell.execute_reply.started": "2023-11-16T20:33:07.688672Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Sequences will be extracted from columns ['utr']\n", "Sequences will be extracted from columns ['utr']\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Calling wandb.login() after wandb.init() has no effect.\n", "\u001b[34m\u001b[1mwandb\u001b[0m: Downloading large artifact human_state_dict:latest, 939.29MB. 1 files... \n", "\u001b[34m\u001b[1mwandb\u001b[0m: 1 of 1 files downloaded. \n", "Done. 0:0:0.8\n" ] } ], "source": [ "# Make labeled datasets\n", "df_train = df_train[['utr', 'rl']]\n", "df_val = df_val[['utr', 'rl']]\n", "train_dataset = grelu.data.dataset.DFSeqDataset(df_train)\n", "val_dataset = grelu.data.dataset.DFSeqDataset(df_val)\n", "\n", "# Build model\n", "model = LightningModel(model_params=model_params, train_params=train_params)\n" ] }, { "cell_type": "code", "execution_count": 9, "id": "84549868-fa1a-49e4-94ce-125294196b97", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (cuda), used: True\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/home/zhaoy183/miniconda3/envs/bioseq/lib/python3.12/site-packages/pytorch_lightning/loggers/wandb.py:396: There is a wandb run already in progress and newly created instances of `WandbLogger` will reuse this run. If this is not desired, call `wandb.finish()` before instantiating `WandbLogger`.\n", "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "da25cf3b03e54120ac0f23198b23ce88", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Validation: | | 0/? [00:00\n", " table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n", " .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n", " .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n", " \n", "

Run history:


epoch▁▁▁▁▂▂▂▂▃▃▃▃▃▃▃▃▄▄▄▄▅▅▅▅▆▆▆▆▆▆▆▆▇▇▇▇████
train_loss_epoch█▅▄▃▃▃▂▂▁▁
train_loss_step▇█▇▅▅▄▄▄▄▄▄▄▃▃▄▄▃▃▂▃▂▃▃▂▃▃▂▃▃▂▃▂▃▂▃▃▂▁▁▂
trainer/global_step▁▁▁▂▂▂▂▂▂▃▃▃▃▃▃▄▄▄▄▄▅▅▅▅▅▅▆▆▆▆▆▇▇▇▇▇▇███
val_loss█▆▄▃▂▂▂▁▁▁▁
val_mse█▆▄▃▂▂▂▁▁▁▁
val_pearson▁▅▇████████

Run summary:


epoch9
train_loss_epoch0.41158
train_loss_step0.528
trainer/global_step10189
val_loss0.65718
val_mse0.65787
val_pearson0.62027

" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ " View run train at: https://wandb.ai/sarosavo/UTR-design/runs/zy01qfcw
View project at: https://wandb.ai/sarosavo/UTR-design
Synced 7 W&B file(s), 0 media file(s), 1 artifact file(s) and 1 other file(s)" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "Find logs at: ./wandb/run-20240718_142535-zy01qfcw/logs" ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "text/html": [ "The new W&B backend becomes opt-out in version 0.18.0; try it out with `wandb.require(\"core\")`! See https://wandb.me/wandb-core for more information." ], "text/plain": [ "" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# Train model\n", "trainer = model.train_on_dataset(train_dataset, val_dataset)\n", "wandb.finish()" ] }, { "cell_type": "markdown", "id": "d7ef0eed-3f40-45c9-80cd-e19596e2bf2c", "metadata": { "execution": { "iopub.execute_input": "2023-06-27T22:34:32.409613Z", "iopub.status.busy": "2023-06-27T22:34:32.409135Z", "iopub.status.idle": "2023-06-27T22:34:32.412438Z", "shell.execute_reply": "2023-06-27T22:34:32.411880Z", "shell.execute_reply.started": "2023-06-27T22:34:32.409595Z" }, "tags": [] }, "source": [ "***" ] }, { "cell_type": "markdown", "id": "bc180fcd-2123-4b33-ae12-5b0719744e1c", "metadata": { "tags": [] }, "source": [ "## 3- Test model" ] }, { "cell_type": "code", "execution_count": 10, "id": "002e87d3-e7fb-4a92-8dc6-d96d809045e2", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:29:26.146429Z", "iopub.status.busy": "2023-11-16T20:29:26.145900Z", "iopub.status.idle": "2023-11-16T20:29:30.158161Z", "shell.execute_reply": "2023-11-16T20:29:30.157469Z", "shell.execute_reply.started": "2023-11-16T20:29:26.146405Z" }, "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Sequences will be extracted from columns ['utr']\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (cuda), used: True\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n", "/home/zhaoy183/miniconda3/envs/bioseq/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py:75: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `pytorch_lightning` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default\n", "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "ee7901697b824d04a34510bf03b68bd3", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Testing: | | 0/? [00:00\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "
test_msetest_pearson
rl0.6281180.62487
\n", "" ], "text/plain": [ " test_mse test_pearson\n", "rl 0.628118 0.62487" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "test_dataset = grelu.data.dataset.DFSeqDataset(df_test)\n", "model.test_on_dataset(test_dataset, devices=[1], num_workers=4)" ] }, { "cell_type": "markdown", "id": "126d4024-6bf7-48d4-ac14-67bdfa89e4e6", "metadata": {}, "source": [ "### Correlation plots for each cell type" ] }, { "cell_type": "code", "execution_count": 11, "id": "bfa4964c-1a0a-404b-ba5b-fe0fa10a0d57", "metadata": { "execution": { "iopub.execute_input": "2023-11-16T20:29:30.159553Z", "iopub.status.busy": "2023-11-16T20:29:30.159389Z", "iopub.status.idle": "2023-11-16T20:29:32.787022Z", "shell.execute_reply": "2023-11-16T20:29:32.786270Z", "shell.execute_reply.started": "2023-11-16T20:29:30.159536Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "GPU available: True (cuda), used: True\n", "TPU available: False, using: 0 TPU cores\n", "HPU available: False, using: 0 HPUs\n", "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/home/zhaoy183/miniconda3/envs/bioseq/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:424: The 'predict_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=255` in the `DataLoader` to improve performance.\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "fd16b971dfed4541997a8d7dbcffa852", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Predicting: | | 0/? [00:00