enformer/enformer-training.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rb_ShvB9E8yM"
      },
      "source": [
        "Copyright 2021 DeepMind Technologies Limited\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "     https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kXQjDxgdwUmW"
      },
      "source": [
        "This colab showcases training of the Enformer model published in\n",
        "\n",
        "**\"Effective gene expression prediction from sequence by integrating long-range interactions\"**\n",
        "\n",
        "Žiga Avsec, Vikram Agarwal, Daniel Visentin, Joseph R. Ledsam, Agnieszka Grabska-Barwinska, Kyle R. Taylor, Yannis Assael, John Jumper, Pushmeet Kohli, David R. Kelley\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2AVkKjy3bh_A"
      },
      "source": [
        "## Steps\n",
        "\n",
        "- Setup tf.data.Dataset by directly accessing the Basenji2 data on GCS: `gs://basenji_barnyard/data`\n",
        "- Train the model for a few steps, alternating training on human and mouse data batches\n",
        "- Evaluate the model on human and mouse genomes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sM_PMOT-2Xhi"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NqR7ol3rxrtM"
      },
      "source": [
        "**Start the colab kernel with GPU**: Runtime -\u003e Change runtime type -\u003e GPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vhjR7StI1tZn"
      },
      "source": [
        "### Install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WiDFm-a41tKW",
        "outputId": "8b889c6e-f113-4664-f2c9-91110808ad92"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting dm-sonnet\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/13/28/9185afffefb655ef1a29f4b84aa9f656826408ca2d1b9ffeba81fbfd40ec/dm_sonnet-2.0.0-py3-none-any.whl (254kB)\n",
            "\r\u001b[K     |█▎                              | 10kB 13.3MB/s eta 0:00:01\r\u001b[K     |██▋                             | 20kB 11.7MB/s eta 0:00:01\r\u001b[K     |███▉                            | 30kB 8.7MB/s eta 0:00:01\r\u001b[K     |█████▏                          | 40kB 7.7MB/s eta 0:00:01\r\u001b[K     |██████▍                         | 51kB 4.5MB/s eta 0:00:01\r\u001b[K     |███████▊                        | 61kB 5.1MB/s eta 0:00:01\r\u001b[K     |█████████                       | 71kB 5.1MB/s eta 0:00:01\r\u001b[K     |██████████▎                     | 81kB 5.5MB/s eta 0:00:01\r\u001b[K     |███████████▋                    | 92kB 5.5MB/s eta 0:00:01\r\u001b[K     |████████████▉                   | 102kB 5.7MB/s eta 0:00:01\r\u001b[K     |██████████████▏                 | 112kB 5.7MB/s eta 0:00:01\r\u001b[K     |███████████████▌                | 122kB 5.7MB/s eta 0:00:01\r\u001b[K     |████████████████▊               | 133kB 5.7MB/s eta 0:00:01\r\u001b[K     |██████████████████              | 143kB 5.7MB/s eta 0:00:01\r\u001b[K     |███████████████████▎            | 153kB 5.7MB/s eta 0:00:01\r\u001b[K     |████████████████████▋           | 163kB 5.7MB/s eta 0:00:01\r\u001b[K     |██████████████████████          | 174kB 5.7MB/s eta 0:00:01\r\u001b[K     |███████████████████████▏        | 184kB 5.7MB/s eta 0:00:01\r\u001b[K     |████████████████████████▌       | 194kB 5.7MB/s eta 0:00:01\r\u001b[K     |█████████████████████████▊      | 204kB 5.7MB/s eta 0:00:01\r\u001b[K     |███████████████████████████     | 215kB 5.7MB/s eta 0:00:01\r\u001b[K     |████████████████████████████▎   | 225kB 5.7MB/s eta 0:00:01\r\u001b[K     |█████████████████████████████▋  | 235kB 5.7MB/s eta 0:00:01\r\u001b[K     |███████████████████████████████ | 245kB 5.7MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 256kB 5.7MB/s \n",
            "\u001b[?25hRequirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (4.41.1)\n",
            "Requirement already satisfied: six\u003e=1.12.0 in /usr/local/lib/python3.7/dist-packages (from dm-sonnet) (1.15.0)\n",
            "Requirement already satisfied: absl-py\u003e=0.7.1 in /usr/local/lib/python3.7/dist-packages (from dm-sonnet) (0.12.0)\n",
            "Requirement already satisfied: numpy\u003e=1.16.3 in /usr/local/lib/python3.7/dist-packages (from dm-sonnet) (1.19.5)\n",
            "Requirement already satisfied: wrapt\u003e=1.11.1 in /usr/local/lib/python3.7/dist-packages (from dm-sonnet) (1.12.1)\n",
            "Requirement already satisfied: tabulate\u003e=0.7.5 in /usr/local/lib/python3.7/dist-packages (from dm-sonnet) (0.8.9)\n",
            "Requirement already satisfied: dm-tree\u003e=0.1.1 in /usr/local/lib/python3.7/dist-packages (from dm-sonnet) (0.1.5)\n",
            "Installing collected packages: dm-sonnet\n",
            "Successfully installed dm-sonnet-2.0.0\n"
          ]
        }
      ],
      "source": [
        "!pip install dm-sonnet tqdm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CokqDsb-fxme"
      },
      "outputs": [],
      "source": [
        "# Get enformer source code\n",
        "!wget -q https://raw.githubusercontent.com/deepmind/deepmind-research/master/enformer/attention_module.py\n",
        "!wget -q https://raw.githubusercontent.com/deepmind/deepmind-research/master/enformer/enformer.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xmffZS_306eb"
      },
      "source": [
        "### Import"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hTGOLrbZxNHK",
        "outputId": "f58b5c21-0764-4003-c794-aa89e5d336cc"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "env: TF_ENABLE_GPU_GARBAGE_COLLECTION=false\n"
          ]
        }
      ],
      "source": [
        "import tensorflow as tf\n",
        "# Make sure the GPU is enabled \n",
        "assert tf.config.list_physical_devices('GPU'), 'Start the colab kernel with GPU: Runtime -\u003e Change runtime type -\u003e GPU'\n",
        "\n",
        "# Easier debugging of OOM\n",
        "%env TF_ENABLE_GPU_GARBAGE_COLLECTION=false"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S9ywsUmT05C1"
      },
      "outputs": [],
      "source": [
        "import sonnet as snt\n",
        "from tqdm import tqdm\n",
        "from IPython.display import clear_output\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import time\n",
        "import os"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YUIbu0Xu1BnA"
      },
      "outputs": [],
      "source": [
        "assert snt.__version__.startswith('2.0')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "id": "PWzsyJddILcx",
        "outputId": "3f1cac0f-6bce-430e-b3c0-9848d43e654c"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            },
            "text/plain": [
              "'2.4.1'"
            ]
          },
          "execution_count": 7,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tf.__version__"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xOhdaXG95eOl",
        "outputId": "1e57ef49-254a-4050-89af-61bc0f8ea577"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Fri Mar 26 12:28:00 2021       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.67       Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   42C    P8    10W /  70W |      3MiB / 15109MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "# GPU colab has T4 with 16 GiB of memory\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Xx--Nco09fN"
      },
      "source": [
        "### Code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BbXyDdoShFzX"
      },
      "outputs": [],
      "source": [
        "import enformer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MEb8OZli2Nbu"
      },
      "outputs": [],
      "source": [
        "# @title `get_targets(organism)`\n",
        "def get_targets(organism):\n",
        "  targets_txt = f'https://raw.githubusercontent.com/calico/basenji/master/manuscripts/cross2020/targets_{organism}.txt'\n",
        "  return pd.read_csv(targets_txt, sep='\\t')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2BuZ2gmUbpXZ"
      },
      "outputs": [],
      "source": [
        "# @title `get_dataset(organism, subset, num_threads=8)`\n",
        "import glob\n",
        "import json\n",
        "import functools\n",
        "\n",
        "\n",
        "def organism_path(organism):\n",
        "  return os.path.join('gs://basenji_barnyard/data', organism)\n",
        "\n",
        "\n",
        "def get_dataset(organism, subset, num_threads=8):\n",
        "  metadata = get_metadata(organism)\n",
        "  dataset = tf.data.TFRecordDataset(tfrecord_files(organism, subset),\n",
        "                                    compression_type='ZLIB',\n",
        "                                    num_parallel_reads=num_threads)\n",
        "  dataset = dataset.map(functools.partial(deserialize, metadata=metadata),\n",
        "                        num_parallel_calls=num_threads)\n",
        "  return dataset\n",
        "\n",
        "\n",
        "def get_metadata(organism):\n",
        "  # Keys:\n",
        "  # num_targets, train_seqs, valid_seqs, test_seqs, seq_length,\n",
        "  # pool_width, crop_bp, target_length\n",
        "  path = os.path.join(organism_path(organism), 'statistics.json')\n",
        "  with tf.io.gfile.GFile(path, 'r') as f:\n",
        "    return json.load(f)\n",
        "\n",
        "\n",
        "def tfrecord_files(organism, subset):\n",
        "  # Sort the values by int(*).\n",
        "  return sorted(tf.io.gfile.glob(os.path.join(\n",
        "      organism_path(organism), 'tfrecords', f'{subset}-*.tfr'\n",
        "  )), key=lambda x: int(x.split('-')[-1].split('.')[0]))\n",
        "\n",
        "\n",
        "def deserialize(serialized_example, metadata):\n",
        "  \"\"\"Deserialize bytes stored in TFRecordFile.\"\"\"\n",
        "  feature_map = {\n",
        "      'sequence': tf.io.FixedLenFeature([], tf.string),\n",
        "      'target': tf.io.FixedLenFeature([], tf.string),\n",
        "  }\n",
        "  example = tf.io.parse_example(serialized_example, feature_map)\n",
        "  sequence = tf.io.decode_raw(example['sequence'], tf.bool)\n",
        "  sequence = tf.reshape(sequence, (metadata['seq_length'], 4))\n",
        "  sequence = tf.cast(sequence, tf.float32)\n",
        "\n",
        "  target = tf.io.decode_raw(example['target'], tf.float16)\n",
        "  target = tf.reshape(target,\n",
        "                      (metadata['target_length'], metadata['num_targets']))\n",
        "  target = tf.cast(target, tf.float32)\n",
        "\n",
        "  return {'sequence': sequence,\n",
        "          'target': target}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VzGRXfwV4tYH"
      },
      "source": [
        "## Load dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 203
        },
        "id": "M_vr1mbl3jbD",
        "outputId": "2de351ed-f43e-4469-a681-2a437d97c946"
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "\u003cdiv\u003e\n",
              "\u003cstyle scoped\u003e\n",
              "    .dataframe tbody tr th:only-of-type {\n",
              "        vertical-align: middle;\n",
              "    }\n",
              "\n",
              "    .dataframe tbody tr th {\n",
              "        vertical-align: top;\n",
              "    }\n",
              "\n",
              "    .dataframe thead th {\n",
              "        text-align: right;\n",
              "    }\n",
              "\u003c/style\u003e\n",
              "\u003ctable border=\"1\" class=\"dataframe\"\u003e\n",
              "  \u003cthead\u003e\n",
              "    \u003ctr style=\"text-align: right;\"\u003e\n",
              "      \u003cth\u003e\u003c/th\u003e\n",
              "      \u003cth\u003eindex\u003c/th\u003e\n",
              "      \u003cth\u003egenome\u003c/th\u003e\n",
              "      \u003cth\u003eidentifier\u003c/th\u003e\n",
              "      \u003cth\u003efile\u003c/th\u003e\n",
              "      \u003cth\u003eclip\u003c/th\u003e\n",
              "      \u003cth\u003escale\u003c/th\u003e\n",
              "      \u003cth\u003esum_stat\u003c/th\u003e\n",
              "      \u003cth\u003edescription\u003c/th\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/thead\u003e\n",
              "  \u003ctbody\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e0\u003c/th\u003e\n",
              "      \u003ctd\u003e0\u003c/td\u003e\n",
              "      \u003ctd\u003e0\u003c/td\u003e\n",
              "      \u003ctd\u003eENCFF833POA\u003c/td\u003e\n",
              "      \u003ctd\u003e/home/drk/tillage/datasets/human/dnase/encode/...\u003c/td\u003e\n",
              "      \u003ctd\u003e32\u003c/td\u003e\n",
              "      \u003ctd\u003e2\u003c/td\u003e\n",
              "      \u003ctd\u003emean\u003c/td\u003e\n",
              "      \u003ctd\u003eDNASE:cerebellum male adult (27 years) and mal...\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e1\u003c/th\u003e\n",
              "      \u003ctd\u003e1\u003c/td\u003e\n",
              "      \u003ctd\u003e0\u003c/td\u003e\n",
              "      \u003ctd\u003eENCFF110QGM\u003c/td\u003e\n",
              "      \u003ctd\u003e/home/drk/tillage/datasets/human/dnase/encode/...\u003c/td\u003e\n",
              "      \u003ctd\u003e32\u003c/td\u003e\n",
              "      \u003ctd\u003e2\u003c/td\u003e\n",
              "      \u003ctd\u003emean\u003c/td\u003e\n",
              "      \u003ctd\u003eDNASE:frontal cortex male adult (27 years) and...\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e2\u003c/th\u003e\n",
              "      \u003ctd\u003e2\u003c/td\u003e\n",
              "      \u003ctd\u003e0\u003c/td\u003e\n",
              "      \u003ctd\u003eENCFF880MKD\u003c/td\u003e\n",
              "      \u003ctd\u003e/home/drk/tillage/datasets/human/dnase/encode/...\u003c/td\u003e\n",
              "      \u003ctd\u003e32\u003c/td\u003e\n",
              "      \u003ctd\u003e2\u003c/td\u003e\n",
              "      \u003ctd\u003emean\u003c/td\u003e\n",
              "      \u003ctd\u003eDNASE:chorion\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e3\u003c/th\u003e\n",
              "      \u003ctd\u003e3\u003c/td\u003e\n",
              "      \u003ctd\u003e0\u003c/td\u003e\n",
              "      \u003ctd\u003eENCFF463ZLQ\u003c/td\u003e\n",
              "      \u003ctd\u003e/home/drk/tillage/datasets/human/dnase/encode/...\u003c/td\u003e\n",
              "      \u003ctd\u003e32\u003c/td\u003e\n",
              "      \u003ctd\u003e2\u003c/td\u003e\n",
              "      \u003ctd\u003emean\u003c/td\u003e\n",
              "      \u003ctd\u003eDNASE:Ishikawa treated with 0.02% dimethyl sul...\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "    \u003ctr\u003e\n",
              "      \u003cth\u003e4\u003c/th\u003e\n",
              "      \u003ctd\u003e4\u003c/td\u003e\n",
              "      \u003ctd\u003e0\u003c/td\u003e\n",
              "      \u003ctd\u003eENCFF890OGQ\u003c/td\u003e\n",
              "      \u003ctd\u003e/home/drk/tillage/datasets/human/dnase/encode/...\u003c/td\u003e\n",
              "      \u003ctd\u003e32\u003c/td\u003e\n",
              "      \u003ctd\u003e2\u003c/td\u003e\n",
              "      \u003ctd\u003emean\u003c/td\u003e\n",
              "      \u003ctd\u003eDNASE:GM03348\u003c/td\u003e\n",
              "    \u003c/tr\u003e\n",
              "  \u003c/tbody\u003e\n",
              "\u003c/table\u003e\n",
              "\u003c/div\u003e"
            ],
            "text/plain": [
              "   index  genome  ... sum_stat                                        description\n",
              "0      0       0  ...     mean  DNASE:cerebellum male adult (27 years) and mal...\n",
              "1      1       0  ...     mean  DNASE:frontal cortex male adult (27 years) and...\n",
              "2      2       0  ...     mean                                      DNASE:chorion\n",
              "3      3       0  ...     mean  DNASE:Ishikawa treated with 0.02% dimethyl sul...\n",
              "4      4       0  ...     mean                                      DNASE:GM03348\n",
              "\n",
              "[5 rows x 8 columns]"
            ]
          },
          "execution_count": 43,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "df_targets_human = get_targets('human')\n",
        "df_targets_human.head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YDSKttXI4hMT"
      },
      "outputs": [],
      "source": [
        "human_dataset = get_dataset('human', 'train').batch(1).repeat()\n",
        "mouse_dataset = get_dataset('mouse', 'train').batch(1).repeat()\n",
        "human_mouse_dataset = tf.data.Dataset.zip((human_dataset, mouse_dataset)).prefetch(2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2vx3116C7oFW"
      },
      "outputs": [],
      "source": [
        "it = iter(mouse_dataset)\n",
        "example = next(it)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XeztqJZ74ixT",
        "outputId": "39dc4051-5a19-4443-b6b0-bf6869faf5ec"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "human\n",
            "{'sequence': (TensorShape([1, 131072, 4]), tf.float32), 'target': (TensorShape([1, 896, 5313]), tf.float32)}\n",
            "mouse\n",
            "{'sequence': (TensorShape([1, 131072, 4]), tf.float32), 'target': (TensorShape([1, 896, 1643]), tf.float32)}\n"
          ]
        }
      ],
      "source": [
        "# Example input\n",
        "it = iter(human_mouse_dataset)\n",
        "example = next(it)\n",
        "for i in range(len(example)):\n",
        "  print(['human', 'mouse'][i])\n",
        "  print({k: (v.shape, v.dtype) for k,v in example[i].items()})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SHHNHzFVbvTk"
      },
      "source": [
        "## Model training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0U3hLJaUdZkG"
      },
      "outputs": [],
      "source": [
        "def create_step_function(model, optimizer):\n",
        "\n",
        "  @tf.function\n",
        "  def train_step(batch, head, optimizer_clip_norm_global=0.2):\n",
        "    with tf.GradientTape() as tape:\n",
        "      outputs = model(batch['sequence'], is_training=True)[head]\n",
        "      loss = tf.reduce_mean(\n",
        "          tf.keras.losses.poisson(batch['target'], outputs))\n",
        "\n",
        "    gradients = tape.gradient(loss, model.trainable_variables)\n",
        "    optimizer.apply(gradients, model.trainable_variables)\n",
        "\n",
        "    return loss\n",
        "  return train_step"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZXv5HU_242Ut"
      },
      "outputs": [],
      "source": [
        "learning_rate = tf.Variable(0., trainable=False, name='learning_rate')\n",
        "optimizer = snt.optimizers.Adam(learning_rate=learning_rate)\n",
        "num_warmup_steps = 5000\n",
        "target_learning_rate = 0.0005\n",
        "\n",
        "model = enformer.Enformer(channels=1536 // 4,  # Use 4x fewer channels to train faster.\n",
        "                          num_heads=8,\n",
        "                          num_transformer_layers=11,\n",
        "                          pooling_type='max')\n",
        "\n",
        "train_step = create_step_function(model, optimizer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "code",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FrbDaOMWcFUl",
        "outputId": "6a42f69c-3003-47f2-a8d2-1b94c52eb57e"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 20/20 [00:24\u003c00:00,  1.25s/it]\n",
            "  0%|          | 0/20 [00:00\u003c?, ?it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "loss_human 1.774059 loss_mouse 0.94303024 learning_rate 2.0000002e-06\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 20/20 [00:17\u003c00:00,  1.13it/s]\n",
            "  0%|          | 0/20 [00:00\u003c?, ?it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "loss_human 1.0067647 loss_mouse 0.8752468 learning_rate 4.0000004e-06\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 20/20 [00:17\u003c00:00,  1.13it/s]\n",
            "  0%|          | 0/20 [00:00\u003c?, ?it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "loss_human 1.0471998 loss_mouse 0.89318746 learning_rate 6e-06\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 20/20 [00:17\u003c00:00,  1.14it/s]\n",
            "  0%|          | 0/20 [00:00\u003c?, ?it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "loss_human 1.010262 loss_mouse 1.02991 learning_rate 8.000001e-06\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 20/20 [00:17\u003c00:00,  1.14it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "loss_human 1.111991 loss_mouse 0.84773445 learning_rate 1.0000001e-05\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "# Train the model\n",
        "steps_per_epoch = 20\n",
        "num_epochs = 5\n",
        "\n",
        "data_it = iter(human_mouse_dataset)\n",
        "global_step = 0\n",
        "for epoch_i in range(num_epochs):\n",
        "  for i in tqdm(range(steps_per_epoch)):\n",
        "    global_step += 1\n",
        "\n",
        "    if global_step \u003e 1:\n",
        "      learning_rate_frac = tf.math.minimum(\n",
        "          1.0, global_step / tf.math.maximum(1.0, num_warmup_steps))      \n",
        "      learning_rate.assign(target_learning_rate * learning_rate_frac)\n",
        "\n",
        "    batch_human, batch_mouse = next(data_it)\n",
        "\n",
        "    loss_human = train_step(batch=batch_human, head='human')\n",
        "    loss_mouse = train_step(batch=batch_mouse, head='mouse')\n",
        "\n",
        "  # End of epoch.\n",
        "  print('')\n",
        "  print('loss_human', loss_human.numpy(),\n",
        "        'loss_mouse', loss_mouse.numpy(),\n",
        "        'learning_rate', optimizer.learning_rate.numpy()\n",
        "        )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cs0f0z0RcCfz"
      },
      "source": [
        "## Evaluate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8c4lNQrHkXSC"
      },
      "outputs": [],
      "source": [
        "# @title `PearsonR` and `R2` metrics\n",
        "\n",
        "def _reduced_shape(shape, axis):\n",
        "  if axis is None:\n",
        "    return tf.TensorShape([])\n",
        "  return tf.TensorShape([d for i, d in enumerate(shape) if i not in axis])\n",
        "\n",
        "\n",
        "class CorrelationStats(tf.keras.metrics.Metric):\n",
        "  \"\"\"Contains shared code for PearsonR and R2.\"\"\"\n",
        "\n",
        "  def __init__(self, reduce_axis=None, name='pearsonr'):\n",
        "    \"\"\"Pearson correlation coefficient.\n",
        "\n",
        "    Args:\n",
        "      reduce_axis: Specifies over which axis to compute the correlation (say\n",
        "        (0, 1). If not specified, it will compute the correlation across the\n",
        "        whole tensor.\n",
        "      name: Metric name.\n",
        "    \"\"\"\n",
        "    super(CorrelationStats, self).__init__(name=name)\n",
        "    self._reduce_axis = reduce_axis\n",
        "    self._shape = None  # Specified in _initialize.\n",
        "\n",
        "  def _initialize(self, input_shape):\n",
        "    # Remaining dimensions after reducing over self._reduce_axis.\n",
        "    self._shape = _reduced_shape(input_shape, self._reduce_axis)\n",
        "\n",
        "    weight_kwargs = dict(shape=self._shape, initializer='zeros')\n",
        "    self._count = self.add_weight(name='count', **weight_kwargs)\n",
        "    self._product_sum = self.add_weight(name='product_sum', **weight_kwargs)\n",
        "    self._true_sum = self.add_weight(name='true_sum', **weight_kwargs)\n",
        "    self._true_squared_sum = self.add_weight(name='true_squared_sum',\n",
        "                                             **weight_kwargs)\n",
        "    self._pred_sum = self.add_weight(name='pred_sum', **weight_kwargs)\n",
        "    self._pred_squared_sum = self.add_weight(name='pred_squared_sum',\n",
        "                                             **weight_kwargs)\n",
        "\n",
        "  def update_state(self, y_true, y_pred, sample_weight=None):\n",
        "    \"\"\"Update the metric state.\n",
        "\n",
        "    Args:\n",
        "      y_true: Multi-dimensional float tensor [batch, ...] containing the ground\n",
        "        truth values.\n",
        "      y_pred: float tensor with the same shape as y_true containing predicted\n",
        "        values.\n",
        "      sample_weight: 1D tensor aligned with y_true batch dimension specifying\n",
        "        the weight of individual observations.\n",
        "    \"\"\"\n",
        "    if self._shape is None:\n",
        "      # Explicit initialization check.\n",
        "      self._initialize(y_true.shape)\n",
        "    y_true.shape.assert_is_compatible_with(y_pred.shape)\n",
        "    y_true = tf.cast(y_true, 'float32')\n",
        "    y_pred = tf.cast(y_pred, 'float32')\n",
        "\n",
        "    self._product_sum.assign_add(\n",
        "        tf.reduce_sum(y_true * y_pred, axis=self._reduce_axis))\n",
        "\n",
        "    self._true_sum.assign_add(\n",
        "        tf.reduce_sum(y_true, axis=self._reduce_axis))\n",
        "\n",
        "    self._true_squared_sum.assign_add(\n",
        "        tf.reduce_sum(tf.math.square(y_true), axis=self._reduce_axis))\n",
        "\n",
        "    self._pred_sum.assign_add(\n",
        "        tf.reduce_sum(y_pred, axis=self._reduce_axis))\n",
        "\n",
        "    self._pred_squared_sum.assign_add(\n",
        "        tf.reduce_sum(tf.math.square(y_pred), axis=self._reduce_axis))\n",
        "\n",
        "    self._count.assign_add(\n",
        "        tf.reduce_sum(tf.ones_like(y_true), axis=self._reduce_axis))\n",
        "\n",
        "  def result(self):\n",
        "    raise NotImplementedError('Must be implemented in subclasses.')\n",
        "\n",
        "  def reset_states(self):\n",
        "    if self._shape is not None:\n",
        "      tf.keras.backend.batch_set_value([(v, np.zeros(self._shape))\n",
        "                                        for v in self.variables])\n",
        "\n",
        "\n",
        "class PearsonR(CorrelationStats):\n",
        "  \"\"\"Pearson correlation coefficient.\n",
        "\n",
        "  Computed as:\n",
        "  ((x - x_avg) * (y - y_avg) / sqrt(Var[x] * Var[y])\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, reduce_axis=(0,), name='pearsonr'):\n",
        "    \"\"\"Pearson correlation coefficient.\n",
        "\n",
        "    Args:\n",
        "      reduce_axis: Specifies over which axis to compute the correlation.\n",
        "      name: Metric name.\n",
        "    \"\"\"\n",
        "    super(PearsonR, self).__init__(reduce_axis=reduce_axis,\n",
        "                                   name=name)\n",
        "\n",
        "  def result(self):\n",
        "    true_mean = self._true_sum / self._count\n",
        "    pred_mean = self._pred_sum / self._count\n",
        "\n",
        "    covariance = (self._product_sum\n",
        "                  - true_mean * self._pred_sum\n",
        "                  - pred_mean * self._true_sum\n",
        "                  + self._count * true_mean * pred_mean)\n",
        "\n",
        "    true_var = self._true_squared_sum - self._count * tf.math.square(true_mean)\n",
        "    pred_var = self._pred_squared_sum - self._count * tf.math.square(pred_mean)\n",
        "    tp_var = tf.math.sqrt(true_var) * tf.math.sqrt(pred_var)\n",
        "    correlation = covariance / tp_var\n",
        "\n",
        "    return correlation\n",
        "\n",
        "\n",
        "class R2(CorrelationStats):\n",
        "  \"\"\"R-squared  (fraction of explained variance).\"\"\"\n",
        "\n",
        "  def __init__(self, reduce_axis=None, name='R2'):\n",
        "    \"\"\"R-squared metric.\n",
        "\n",
        "    Args:\n",
        "      reduce_axis: Specifies over which axis to compute the correlation.\n",
        "      name: Metric name.\n",
        "    \"\"\"\n",
        "    super(R2, self).__init__(reduce_axis=reduce_axis,\n",
        "                             name=name)\n",
        "\n",
        "  def result(self):\n",
        "    true_mean = self._true_sum / self._count\n",
        "    total = self._true_squared_sum - self._count * tf.math.square(true_mean)\n",
        "    residuals = (self._pred_squared_sum - 2 * self._product_sum\n",
        "                 + self._true_squared_sum)\n",
        "\n",
        "    return tf.ones_like(residuals) - residuals / total\n",
        "\n",
        "\n",
        "class MetricDict:\n",
        "  def __init__(self, metrics):\n",
        "    self._metrics = metrics\n",
        "\n",
        "  def update_state(self, y_true, y_pred):\n",
        "    for k, metric in self._metrics.items():\n",
        "      metric.update_state(y_true, y_pred)\n",
        "\n",
        "  def result(self):\n",
        "    return {k: metric.result() for k, metric in self._metrics.items()}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x80gX9LrhBU-"
      },
      "outputs": [],
      "source": [
        "def evaluate_model(model, dataset, head, max_steps=None):\n",
        "  metric = MetricDict({'PearsonR': PearsonR(reduce_axis=(0,1))})\n",
        "  @tf.function\n",
        "  def predict(x):\n",
        "    return model(x, is_training=False)[head]\n",
        "\n",
        "  for i, batch in tqdm(enumerate(dataset)):\n",
        "    if max_steps is not None and i \u003e max_steps:\n",
        "      break\n",
        "    metric.update_state(batch['target'], predict(batch['sequence']))\n",
        "\n",
        "  return metric.result()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "57fNitK9hzwd",
        "outputId": "947aaadb-dad2-4a00-ddac-d765f65d782f"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "101it [00:23,  6.27it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "{'PearsonR': 0.0028573992}\n"
          ]
        }
      ],
      "source": [
        "metrics_human = evaluate_model(model,\n",
        "                               dataset=get_dataset('human', 'valid').batch(1).prefetch(2),\n",
        "                               head='human',\n",
        "                               max_steps=100)\n",
        "print('')\n",
        "print({k: v.numpy().mean() for k, v in metrics_human.items()})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HY_wj95xiDtE",
        "outputId": "fea839f7-b6c9-46ed-aece-c56b02e9ea16"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "101it [00:21,  6.54it/s]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "{'PearsonR': 0.005183698}\n"
          ]
        }
      ],
      "source": [
        "metrics_mouse = evaluate_model(model,\n",
        "                               dataset=get_dataset('mouse', 'valid').batch(1).prefetch(2),\n",
        "                               head='mouse',\n",
        "                               max_steps=100)\n",
        "print('')\n",
        "print({k: v.numpy().mean() for k, v in metrics_mouse.items()})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5k1yaJrNCgvw"
      },
      "source": [
        "# Restore Checkpoint\n",
        "\n",
        "Note: For the TF-Hub Enformer model, the required input sequence length is 393,216 which actually gets cropped within the model to 196,608. The open source module does not internally crop the sequence. Therefore, the code below crops the central `196,608 bp` of the longer sequence to reproduce the output of the TF hub from the reloaded checkpoint."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DB2cGdH8EGfn"
      },
      "outputs": [],
      "source": [
        "np.random.seed(42)\n",
        "EXTENDED_SEQ_LENGTH = 393_216\n",
        "SEQ_LENGTH = 196_608\n",
        "inputs = np.array(np.random.random((1, EXTENDED_SEQ_LENGTH, 4)), dtype=np.float32)\n",
        "inputs_cropped = enformer.TargetLengthCrop1D(SEQ_LENGTH)(inputs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mdf35itsCjEY"
      },
      "outputs": [],
      "source": [
        "checkpoint_gs_path = 'gs://dm-enformer/models/enformer/sonnet_weights/*'\n",
        "checkpoint_path = '/tmp/enformer_checkpoint'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 280,
          "status": "ok",
          "timestamp": 1653476327690,
          "user": {
            "displayName": "Kyle Taylor",
            "userId": "14169907681771397124"
          },
          "user_tz": -60
        },
        "id": "G2P4IHqswLul",
        "outputId": "180abe21-ba00-4031-d9d7-2326f1f742f2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "mkdir: cannot create directory ‘/tmp/enformer_checkpoint’: File exists\n"
          ]
        }
      ],
      "source": [
        "!mkdir /tmp/enformer_checkpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 19279,
          "status": "ok",
          "timestamp": 1653476347208,
          "user": {
            "displayName": "Kyle Taylor",
            "userId": "14169907681771397124"
          },
          "user_tz": -60
        },
        "id": "LTL8EISGCujC",
        "outputId": "2b743b9b-480d-44dc-b08e-82d2bc089a47"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "gs://dm-enformer/models/enformer/sonnet_weights/checkpoint\n",
            "gs://dm-enformer/models/enformer/sonnet_weights/enformer-fine-tuned-human-1.data-00000-of-00001\n",
            "gs://dm-enformer/models/enformer/sonnet_weights/enformer-fine-tuned-human-1.index\n"
          ]
        }
      ],
      "source": [
        "# Copy checkpoints from GCS to temporary directory.\n",
        "# This will take a while as the checkpoint is ~ 1GB.\n",
        "for file_path in tf.io.gfile.glob(checkpoint_gs_path):\n",
        "  print(file_path)\n",
        "  file_name = os.path.basename(file_path)\n",
        "  tf.io.gfile.copy(file_path, f'{checkpoint_path}/{file_name}', overwrite=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 322,
          "status": "ok",
          "timestamp": 1653476347527,
          "user": {
            "displayName": "Kyle Taylor",
            "userId": "14169907681771397124"
          },
          "user_tz": -60
        },
        "id": "9VSeTx0sCvcw",
        "outputId": "b52d7570-c355-4068-b932-3796b56e5586"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "total 959M\n",
            "-rw-r--r-- 1 root root  111 May 25 10:58 checkpoint\n",
            "-rw-r--r-- 1 root root 959M May 25 10:59 enformer-fine-tuned-human-1.data-00000-of-00001\n",
            "-rw-r--r-- 1 root root 5.7K May 25 10:59 enformer-fine-tuned-human-1.index\n"
          ]
        }
      ],
      "source": [
        "!ls -lh /tmp/enformer_checkpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "00Y2GgRED3aI"
      },
      "outputs": [],
      "source": [
        "enformer_model = enformer.Enformer()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mFyIiGyiD5yh"
      },
      "outputs": [],
      "source": [
        "checkpoint = tf.train.Checkpoint(module=enformer_model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 6,
          "status": "ok",
          "timestamp": 1653476347529,
          "user": {
            "displayName": "Kyle Taylor",
            "userId": "14169907681771397124"
          },
          "user_tz": -60
        },
        "id": "VuyspnpOD9kA",
        "outputId": "3b495138-be27-4459-af82-c5da2af5bd2d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/tmp/enformer_checkpoint/enformer-fine-tuned-human-1\n"
          ]
        }
      ],
      "source": [
        "latest = tf.train.latest_checkpoint(checkpoint_path)\n",
        "print(latest)\n",
        "status = checkpoint.restore(latest)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MKkVOTyKEABJ"
      },
      "outputs": [],
      "source": [
        "# Using `is_training=False` to match TF-hub predict_on_batch function.\n",
        "restored_predictions = enformer_model(inputs_cropped, is_training=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OX650jqCEQdv"
      },
      "outputs": [],
      "source": [
        "import tensorflow_hub as hub\n",
        "enformer_tf_hub_model = hub.load(\"https://tfhub.dev/deepmind/enformer/1\").model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yiOTFTSdE5H1"
      },
      "outputs": [],
      "source": [
        "hub_predictions = enformer_tf_hub_model.predict_on_batch(inputs)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "executionInfo": {
          "elapsed": 15,
          "status": "ok",
          "timestamp": 1653476357260,
          "user": {
            "displayName": "Kyle Taylor",
            "userId": "14169907681771397124"
          },
          "user_tz": -60
        },
        "id": "uYrWgfaGFbpL",
        "outputId": "cb9d3ad0-ee14-46ac-b188-a4b6a697a159"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "execution_count": 41,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "np.allclose(hub_predictions['human'], restored_predictions['human'], atol=1e-5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4wEOUMeNzK8q"
      },
      "outputs": [],
      "source": [
        "# Can run with 'is_training=True' but note that this will\n",
        "# change the predictions as the batch statistics will be updated\n",
        "# and the outputs will likley not match the TF-hub model.\n",
        "# enformer(inputs_cropped, is_training=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jyVHRPAN5w6J"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/research/genomics/dna_to_rna:colab",
        "kind": "private"
      },
      "name": "enformer-training.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}