ogbg_molhiv.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/PabloExperimental/OGB/blob/main/ogbg_molhiv.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u4S_hgwe-1b4"
      },
      "source": [
        "Original: [Github](https://github.com/danielegrattarola/spektral/blob/master/examples/graph_prediction/ogbg-mol-hiv_ecc.py)\n",
        "\n",
        "[OGB HIV Overview on Hugging Face](https://huggingface.co/datasets/OGB/ogbg-molhiv)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "id": "7gepIlmY_KPj",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c1e9a123-89e0-46e4-836b-b03ae28b4fd6"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: tensorflow==2.15.0 in /usr/local/lib/python3.10/dist-packages (2.15.0)\n",
            "Requirement already satisfied: spektral in /usr/local/lib/python3.10/dist-packages (1.3.1)\n",
            "Requirement already satisfied: ogb in /usr/local/lib/python3.10/dist-packages (1.3.6)\n",
            "Requirement already satisfied: absl-py>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (1.4.0)\n",
            "Requirement already satisfied: astunparse>=1.6.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (1.6.3)\n",
            "Requirement already satisfied: flatbuffers>=23.5.26 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (24.3.25)\n",
            "Requirement already satisfied: gast!=0.5.0,!=0.5.1,!=0.5.2,>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (0.6.0)\n",
            "Requirement already satisfied: google-pasta>=0.1.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (0.2.0)\n",
            "Requirement already satisfied: h5py>=2.9.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (3.11.0)\n",
            "Requirement already satisfied: libclang>=13.0.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (18.1.1)\n",
            "Requirement already satisfied: ml-dtypes~=0.2.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (0.2.0)\n",
            "Requirement already satisfied: numpy<2.0.0,>=1.23.5 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (1.26.4)\n",
            "Requirement already satisfied: opt-einsum>=2.3.2 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (3.4.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (24.1)\n",
            "Requirement already satisfied: protobuf!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5,<5.0.0dev,>=3.20.3 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (3.20.3)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (75.1.0)\n",
            "Requirement already satisfied: six>=1.12.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (1.16.0)\n",
            "Requirement already satisfied: termcolor>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (2.5.0)\n",
            "Requirement already satisfied: typing-extensions>=3.6.6 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (4.12.2)\n",
            "Requirement already satisfied: wrapt<1.15,>=1.11.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (1.14.1)\n",
            "Requirement already satisfied: tensorflow-io-gcs-filesystem>=0.23.1 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (0.37.1)\n",
            "Requirement already satisfied: grpcio<2.0,>=1.24.3 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (1.64.1)\n",
            "Requirement already satisfied: tensorboard<2.16,>=2.15 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (2.15.2)\n",
            "Requirement already satisfied: tensorflow-estimator<2.16,>=2.15.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (2.15.0)\n",
            "Requirement already satisfied: keras<2.16,>=2.15.0 in /usr/local/lib/python3.10/dist-packages (from tensorflow==2.15.0) (2.15.0)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.10/dist-packages (from spektral) (1.4.2)\n",
            "Requirement already satisfied: lxml in /usr/local/lib/python3.10/dist-packages (from spektral) (4.9.4)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from spektral) (3.4.2)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from spektral) (2.2.2)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from spektral) (2.32.3)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.10/dist-packages (from spektral) (1.5.2)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from spektral) (1.13.1)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from spektral) (4.66.5)\n",
            "Requirement already satisfied: torch>=1.6.0 in /usr/local/lib/python3.10/dist-packages (from ogb) (2.5.0+cu121)\n",
            "Requirement already satisfied: urllib3>=1.24.0 in /usr/local/lib/python3.10/dist-packages (from ogb) (2.2.3)\n",
            "Requirement already satisfied: outdated>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from ogb) (0.2.2)\n",
            "Requirement already satisfied: wheel<1.0,>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from astunparse>=1.6.0->tensorflow==2.15.0) (0.44.0)\n",
            "Requirement already satisfied: littleutils in /usr/local/lib/python3.10/dist-packages (from outdated>=0.2.0->ogb) (0.2.4)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas->spektral) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->spektral) (2024.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.10/dist-packages (from pandas->spektral) (2024.2)\n",
            "Requirement already satisfied: threadpoolctl>=3.1.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->spektral) (3.5.0)\n",
            "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.0) (2.27.0)\n",
            "Requirement already satisfied: google-auth-oauthlib<2,>=0.5 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.0) (1.2.1)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.0) (3.7)\n",
            "Requirement already satisfied: tensorboard-data-server<0.8.0,>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.0) (0.7.2)\n",
            "Requirement already satisfied: werkzeug>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from tensorboard<2.16,>=2.15->tensorflow==2.15.0) (3.0.4)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->spektral) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->spektral) (3.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->spektral) (2024.8.30)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch>=1.6.0->ogb) (3.16.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.6.0->ogb) (3.1.4)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch>=1.6.0->ogb) (2024.6.1)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.6.0->ogb) (1.13.1)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy==1.13.1->torch>=1.6.0->ogb) (1.3.0)\n",
            "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow==2.15.0) (5.5.0)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow==2.15.0) (0.4.1)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.10/dist-packages (from google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow==2.15.0) (4.9)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow==2.15.0) (1.3.1)\n",
            "Requirement already satisfied: MarkupSafe>=2.1.1 in /usr/local/lib/python3.10/dist-packages (from werkzeug>=1.0.1->tensorboard<2.16,>=2.15->tensorflow==2.15.0) (3.0.2)\n",
            "Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in /usr/local/lib/python3.10/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard<2.16,>=2.15->tensorflow==2.15.0) (0.6.1)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.10/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<2,>=0.5->tensorboard<2.16,>=2.15->tensorflow==2.15.0) (3.2.2)\n"
          ]
        }
      ],
      "source": [
        "!pip install tensorflow==2.15.0 spektral ogb"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from ogb.graphproppred import Evaluator, GraphPropPredDataset\n",
        "from tensorflow.keras.layers import Dense, Input, Dropout, BatchNormalization\n",
        "from tensorflow.keras.losses import BinaryCrossentropy\n",
        "from tensorflow.keras.models import Model\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "\n",
        "from spektral.data import DisjointLoader\n",
        "from spektral.datasets import OGB\n",
        "from spektral.layers import ECCConv, GlobalAvgPool\n",
        "\n",
        "################################################################################\n",
        "# Config\n",
        "################################################################################\n",
        "learning_rate = 1e-3  # Learning rate\n",
        "epochs = 15  # Number of training epochs\n",
        "batch_size = 16  # Batch size\n",
        "\n",
        "################################################################################\n",
        "# Load data\n",
        "################################################################################\n",
        "dataset_name = \"ogbg-molhiv\"\n",
        "ogb_dataset = GraphPropPredDataset(name=dataset_name)\n",
        "dataset = OGB(ogb_dataset)\n",
        "\n",
        "# Parameters\n",
        "F = dataset.n_node_features  # Dimension of node features\n",
        "S = dataset.n_edge_features  # Dimension of edge features\n",
        "n_out = dataset.n_labels  # Dimension of the target\n",
        "\n",
        "# Train/test split\n",
        "idx = ogb_dataset.get_idx_split()\n",
        "idx_tr, idx_va, idx_te = idx[\"train\"], idx[\"valid\"], idx[\"test\"]\n",
        "dataset_tr = dataset[idx_tr]\n",
        "dataset_va = dataset[idx_va]\n",
        "dataset_te = dataset[idx_te]\n",
        "\n",
        "loader_tr = DisjointLoader(dataset_tr, batch_size=batch_size, epochs=epochs)\n",
        "loader_te = DisjointLoader(dataset_te, batch_size=batch_size, epochs=1)\n",
        "\n",
        "################################################################################\n",
        "# Build model\n",
        "################################################################################\n",
        "X_in = Input(shape=(F,))\n",
        "A_in = Input(shape=(None,), sparse=True)\n",
        "E_in = Input(shape=(S,))\n",
        "I_in = Input(shape=(), dtype=tf.int64)\n",
        "\n",
        "x = ECCConv(64, activation=\"relu\")([X_in, A_in, E_in])\n",
        "\n",
        "x = Dropout(0.25)(x)\n",
        "\n",
        "x = ECCConv(64, activation=\"relu\")([x, A_in, E_in])\n",
        "\n",
        "x = Dropout(0.25)(x)\n",
        "\n",
        "x = ECCConv(64, activation=\"relu\")([x, A_in, E_in])\n",
        "\n",
        "x = Dropout(0.25)(x)\n",
        "\n",
        "x = GlobalAvgPool()([x, I_in])\n",
        "\n",
        "output = Dense(n_out, activation=\"sigmoid\")(x)\n",
        "\n",
        "model = Model(inputs=[X_in, A_in, E_in, I_in], outputs=output)\n",
        "optimizer = Adam(learning_rate)\n",
        "loss_fn = BinaryCrossentropy()\n",
        "\n",
        "model.summary()\n",
        "\n",
        "################################################################################\n",
        "# Fit model\n",
        "################################################################################\n",
        "@tf.function(input_signature=loader_tr.tf_signature(), experimental_relax_shapes=True)\n",
        "def train_step(inputs, target):\n",
        "    with tf.GradientTape() as tape:\n",
        "        predictions = model(inputs, training=True)\n",
        "        loss = loss_fn(target, predictions) + sum(model.losses)\n",
        "    gradients = tape.gradient(loss, model.trainable_variables)\n",
        "    optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
        "    return loss\n",
        "\n",
        "\n",
        "step = loss = 0\n",
        "for batch in loader_tr:\n",
        "    step += 1\n",
        "    loss += train_step(*batch)\n",
        "    if step == loader_tr.steps_per_epoch:\n",
        "        step = 0\n",
        "        print(\"Loss: {:.4f}\".format(loss / loader_tr.steps_per_epoch))\n",
        "        loss = 0\n",
        "\n",
        "################################################################################\n",
        "# Evaluate model\n",
        "################################################################################\n",
        "print(\"Testing model\")\n",
        "evaluator = Evaluator(name=dataset_name)\n",
        "y_true = []\n",
        "y_pred = []\n",
        "for batch in loader_te:\n",
        "    inputs, target = batch\n",
        "    p = model(inputs, training=False)\n",
        "    y_true.append(target)\n",
        "    y_pred.append(p.numpy())\n",
        "\n",
        "y_true = np.vstack(y_true)\n",
        "y_pred = np.vstack(y_pred)\n",
        "model_loss = loss_fn(y_true, y_pred)\n",
        "ogb_score = evaluator.eval({\"y_true\": y_true, \"y_pred\": y_pred})\n",
        "\n",
        "print(\n",
        "    \"Done. Test loss: {:.4f}. ROC-AUC: {:.2f}\".format(model_loss, ogb_score[\"rocauc\"])\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FSqzEdIoKcpq",
        "outputId": "4cbf8d5f-8302-4741-d76f-da27cd27eb23"
      },
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/ogb/graphproppred/dataset.py:67: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
            "  loaded_dict = torch.load(pre_processed_file_path, 'rb')\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model: \"model_18\"\n",
            "__________________________________________________________________________________________________\n",
            " Layer (type)                Output Shape                 Param #   Connected to                  \n",
            "==================================================================================================\n",
            " input_77 (InputLayer)       [(None, 9)]                  0         []                            \n",
            "                                                                                                  \n",
            " input_78 (InputLayer)       [(None, None)]               0         []                            \n",
            "                                                                                                  \n",
            " input_79 (InputLayer)       [(None, 3)]                  0         []                            \n",
            "                                                                                                  \n",
            " ecc_conv_53 (ECCConv)       (None, 64)                   2944      ['input_77[0][0]',            \n",
            "                                                                     'input_78[0][0]',            \n",
            "                                                                     'input_79[0][0]']            \n",
            "                                                                                                  \n",
            " dropout_41 (Dropout)        (None, 64)                   0         ['ecc_conv_53[0][0]']         \n",
            "                                                                                                  \n",
            " ecc_conv_54 (ECCConv)       (None, 64)                   20544     ['dropout_41[0][0]',          \n",
            "                                                                     'input_78[0][0]',            \n",
            "                                                                     'input_79[0][0]']            \n",
            "                                                                                                  \n",
            " dropout_42 (Dropout)        (None, 64)                   0         ['ecc_conv_54[0][0]']         \n",
            "                                                                                                  \n",
            " ecc_conv_55 (ECCConv)       (None, 64)                   20544     ['dropout_42[0][0]',          \n",
            "                                                                     'input_78[0][0]',            \n",
            "                                                                     'input_79[0][0]']            \n",
            "                                                                                                  \n",
            " dropout_43 (Dropout)        (None, 64)                   0         ['ecc_conv_55[0][0]']         \n",
            "                                                                                                  \n",
            " input_80 (InputLayer)       [(None,)]                    0         []                            \n",
            "                                                                                                  \n",
            " global_avg_pool_2 (GlobalA  (None, 64)                   0         ['dropout_43[0][0]',          \n",
            " vgPool)                                                             'input_80[0][0]']            \n",
            "                                                                                                  \n",
            " dense_19 (Dense)            (None, 1)                    65        ['global_avg_pool_2[0][0]']   \n",
            "                                                                                                  \n",
            "==================================================================================================\n",
            "Total params: 44097 (172.25 KB)\n",
            "Trainable params: 44097 (172.25 KB)\n",
            "Non-trainable params: 0 (0.00 Byte)\n",
            "__________________________________________________________________________________________________\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/spektral/data/utils.py:221: UserWarning: you are shuffling a 'OGB' object which is not a subclass of 'Sequence'; `shuffle` is not guaranteed to behave correctly. E.g., non-numpy array/tensor objects with view semantics may contain duplicates after shuffling.\n",
            "  np.random.shuffle(a)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loss: 0.1674\n",
            "Loss: 0.1609\n",
            "Loss: 0.1561\n",
            "Loss: 0.1551\n",
            "Loss: 0.1517\n",
            "Loss: 0.1493\n",
            "Loss: 0.1488\n",
            "Loss: 0.1464\n",
            "Loss: 0.1475\n",
            "Loss: 0.1451\n",
            "Loss: 0.1442\n",
            "Loss: 0.1424\n",
            "Loss: 0.1417\n",
            "Loss: 0.1412\n",
            "Loss: 0.1403\n",
            "Testing model\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/spektral/data/utils.py:221: UserWarning: you are shuffling a 'OGB' object which is not a subclass of 'Sequence'; `shuffle` is not guaranteed to behave correctly. E.g., non-numpy array/tensor objects with view semantics may contain duplicates after shuffling.\n",
            "  np.random.shuffle(a)\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Done. Test loss: 0.1332. ROC-AUC: 0.74\n"
          ]
        }
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyPty806r8zaYQd1LhjSChxp",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}