GNN_based_Screening_Predictor?!.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GNN-based Screening Predictor?!",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/alex-movila/ML-Colab-Tutorials/blob/master/GNN_based_Screening_Predictor%3F!.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aqWEELs5phU6"
      },
      "outputs": [],
      "source": [
        "!pip install ogb\n",
        "!pip install optax"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from ogb.graphproppred import GraphPropPredDataset\n",
        "\n",
        "dataset = GraphPropPredDataset(name='ogbg-molhiv')\n",
        "\n",
        "graph, label = dataset[0]\n",
        "print('Keys are', graph.keys())\n",
        "print('Label is', label)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b-9lnrVMtsFt",
        "outputId": "602e2914-ddd1-4dfb-fc45-801b3b0f824b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading http://snap.stanford.edu/ogb/data/graphproppred/csv_mol_download/hiv.zip\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloaded 0.00 GB: 100%|██████████| 3/3 [00:00<00:00,  6.70it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Extracting dataset/hiv.zip\n",
            "Loading necessary files...\n",
            "This might take a while.\n",
            "Processing graphs...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 41127/41127 [00:00<00:00, 81409.03it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saving...\n",
            "Keys are dict_keys(['edge_index', 'edge_feat', 'node_feat', 'num_nodes'])\n",
            "Label is [0]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(graph['num_nodes'])\n",
        "print(graph['node_feat'].shape)\n",
        "print(graph['edge_feat'].shape)\n",
        "print(graph['edge_index'].shape)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aXgmyTObuqXY",
        "outputId": "6f395c8d-7f8f-4c5c-ab33-6aaf82329eef"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "19\n",
            "(19, 9)\n",
            "(40, 3)\n",
            "(2, 40)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(graph['node_feat'][:5, :])\n",
        "print(graph['edge_feat'][:5, :])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Silnky6PvZQG",
        "outputId": "558636ab-79db-40c5-f830-d83072200f01"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[ 5  0  4  5  3  0  2  0  0]\n",
            " [ 5  0  4  5  2  0  2  0  0]\n",
            " [ 5  0  3  5  0  0  1  0  1]\n",
            " [ 7  0  2  6  0  0  1  0  1]\n",
            " [28  0  4  2  0  0  5  0  1]]\n",
            "[[0 0 0]\n",
            " [0 0 0]\n",
            " [0 0 0]\n",
            " [0 0 0]\n",
            " [1 0 0]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(graph['edge_index'])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eK8vDjljwZ8S",
        "outputId": "0b0936cf-0ad2-4fff-b8d6-62195b3bd23f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[[ 0  1  1  2  2  3  3  4  4  5  5  6  6  7  7  8  6  9  4 10 10 11 11 12\n",
            "  12 13 11 14 14 15 15 16 16 17 15 18  9  2 18  4]\n",
            " [ 1  0  2  1  3  2  4  3  5  4  6  5  7  6  8  7  9  6 10  4 11 10 12 11\n",
            "  13 12 14 11 15 14 16 15 17 16 18 15  2  9  4 18]]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "\n",
        "def convert_edge_index_to_matrix(edge_index, nb_nodes):\n",
        "  adj_mat = np.eye(nb_nodes)\n",
        "  for i in range(edge_index.shape[1]):\n",
        "    adj_mat[edge_index[0, i], edge_index[1, i]] = 1.0\n",
        "  return adj_mat / np.sum(adj_mat, axis=-1, keepdims=True)\n",
        "\n",
        "print(convert_edge_index_to_matrix(graph['edge_index'], graph['num_nodes']))"
      ],
      "metadata": {
        "id": "e0N3g5Erv3Dp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import jax\n",
        "import jax.numpy as jnp\n",
        "\n",
        "@jax.jit\n",
        "def simple_gnn_layer(weights, features, adj_matrix):\n",
        "  latent = jnp.matmul(features, weights)  # (N x H) * (H x H') -> (N x H')\n",
        "  latent = jnp.matmul(adj_matrix, latent)  # (N x N) * (N x H') -> (N x H')\n",
        "  latent = jax.nn.relu(latent)\n",
        "  return latent\n",
        "\n",
        "@jax.jit\n",
        "def network(params, features, adj_matrix):\n",
        "  latent = features\n",
        "  for layer in range(len(params) - 1):\n",
        "    latent = simple_gnn_layer(params[layer], latent, adj_matrix) # (N x H)\n",
        "  g_features = jnp.mean(latent, axis=0)  # (H,)\n",
        "  logits = jnp.matmul(g_features, params[-1])  # (1,)\n",
        "  return logits"
      ],
      "metadata": {
        "id": "Gi3zJtM9xdzj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@jax.jit\n",
        "def binary_cross_entropy(logits, labels):\n",
        "  max_val = jnp.clip(logits, 0, None)\n",
        "  loss = logits - logits * labels + max_val + jnp.log(\n",
        "      jnp.exp(-max_val) + jnp.exp((-logits - max_val)))\n",
        "  return jnp.mean(loss)\n",
        "\n",
        "@jax.jit\n",
        "def _loss(params, features, adj_matrix, labels):\n",
        "  logits = network(params, features, adj_matrix)\n",
        "  return binary_cross_entropy(logits, labels)\n",
        "\n",
        "@jax.jit\n",
        "def accuracy(logits, labels):\n",
        "  return jnp.mean((logits > 0.0) == (labels > 0.5))"
      ],
      "metadata": {
        "id": "JPUP74gwzc4Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import optax\n",
        "\n",
        "split_idx = dataset.get_idx_split()\n",
        "train_idx, val_idx = split_idx[\"train\"], split_idx[\"valid\"]\n",
        "input_dim = graph['node_feat'].shape[1]\n",
        "\n",
        "def train(hidden_dim, nb_layers, epochs, learning_rate):\n",
        "  params = []\n",
        "  params.append(np.random.randn(input_dim, hidden_dim) / np.sqrt(input_dim))\n",
        "  for i in range(nb_layers - 2):\n",
        "    params.append(np.random.randn(hidden_dim, hidden_dim) / np.sqrt(hidden_dim))\n",
        "  params.append(np.random.randn(hidden_dim, 1) / np.sqrt(hidden_dim))\n",
        "\n",
        "  opt = optax.adam(learning_rate=learning_rate)\n",
        "  opt_state = opt.init(params)\n",
        "\n",
        "  @jax.jit\n",
        "  def _step(params, opt_state, features, adj_matrix, labels):\n",
        "    loss, grads = jax.value_and_grad(_loss)(\n",
        "        params, features, adj_matrix, labels)\n",
        "    updates, opt_state = opt.update(grads, opt_state, params)\n",
        "    params = optax.apply_updates(params, updates)\n",
        "    return params, opt_state, loss\n",
        "  \n",
        "  ep = 0\n",
        "  step = 0\n",
        "  while ep < epochs:\n",
        "    for idx in train_idx:\n",
        "      graph, label = dataset[idx]\n",
        "      node_fts = graph['node_feat']\n",
        "      nb_nodes = graph['num_nodes']\n",
        "      adj_mat = convert_edge_index_to_matrix(graph['edge_index'], nb_nodes)\n",
        "      params, opt_state, loss = _step(\n",
        "          params, opt_state, node_fts, adj_mat, label\n",
        "      )\n",
        "      if step % 1000 == 0:\n",
        "        print(f'step = {step} | loss = {loss}')\n",
        "      step += 1\n",
        "    val_preds = []\n",
        "    val_labels = []\n",
        "    for idx in val_idx:\n",
        "      graph, label = dataset[idx]\n",
        "      node_fts = graph['node_feat']\n",
        "      nb_nodes = graph['num_nodes']\n",
        "      adj_mat = convert_edge_index_to_matrix(graph['edge_index'], nb_nodes)\n",
        "      val_preds.append(network(params, node_fts, adj_mat)[0])\n",
        "      val_labels.append(label)\n",
        "    val_accuracy = accuracy(jnp.array(val_preds), jnp.array(val_labels))\n",
        "    print(f'epoch = {ep} | validation accuracy = {val_accuracy}')\n",
        "    ep += 1\n",
        "  return params"
      ],
      "metadata": {
        "id": "Bo6lgCAZ0xqx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "trained_model = train(hidden_dim=32, nb_layers=2, epochs=1, learning_rate=0.001)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "XdS_buL53MM-",
        "outputId": "b5bf2f9c-eba9-48d2-99d1-b0b90eca721b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "step = 0 | loss = 3.228520154953003\n",
            "step = 1000 | loss = 0.00039005279541015625\n",
            "step = 2000 | loss = 1.0528050661087036\n",
            "step = 3000 | loss = 0.017929553985595703\n",
            "step = 4000 | loss = 0.02083730697631836\n",
            "step = 5000 | loss = 0.016472816467285156\n",
            "step = 6000 | loss = 0.010475635528564453\n",
            "step = 7000 | loss = 0.02108168601989746\n",
            "step = 8000 | loss = 0.0829765796661377\n",
            "step = 9000 | loss = 0.002238750457763672\n",
            "step = 10000 | loss = 0.012306690216064453\n",
            "step = 11000 | loss = 0.0037064552307128906\n",
            "step = 12000 | loss = 0.026071548461914062\n",
            "step = 13000 | loss = 0.012278556823730469\n",
            "step = 14000 | loss = 0.02101278305053711\n",
            "step = 15000 | loss = 0.025160551071166992\n",
            "step = 16000 | loss = 0.029111146926879883\n",
            "step = 17000 | loss = 0.07236838340759277\n",
            "step = 18000 | loss = 0.02436351776123047\n",
            "step = 19000 | loss = 0.018687963485717773\n",
            "step = 20000 | loss = 0.06229758262634277\n",
            "step = 21000 | loss = 0.013743877410888672\n",
            "step = 22000 | loss = 0.018528223037719727\n",
            "step = 23000 | loss = 0.0512700080871582\n",
            "step = 24000 | loss = 0.014118671417236328\n",
            "step = 25000 | loss = 0.04196596145629883\n",
            "step = 26000 | loss = 0.007762432098388672\n",
            "step = 27000 | loss = 0.011636734008789062\n",
            "step = 28000 | loss = 0.01963663101196289\n",
            "step = 29000 | loss = 0.02144145965576172\n",
            "step = 30000 | loss = 0.03717207908630371\n",
            "step = 31000 | loss = 0.058652639389038086\n",
            "step = 32000 | loss = 0.047481536865234375\n",
            "epoch = 0 | validation accuracy = 0.9793721437454224\n"
          ]
        }
      ]
    }
  ]
}