models.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zn5U4EE6K86v",
        "outputId": "c003f212-2283-4f74-eb5f-00c0fb35ea4e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "2.2.1\n"
          ]
        }
      ],
      "source": [
        "# Install required packages.\n",
        "import os\n",
        "import torch\n",
        "from torchmetrics import AUROC\n",
        "os.environ['TORCH'] = torch.__version__\n",
        "print(torch.__version__)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "from torch_geometric.datasets import TUDataset\n",
        "dataset = TUDataset(root='data/TUDataset', name='Tox21_AhR_training')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mHSP6-RBOqCE",
        "outputId": "ae211408-a548-40e5-844e-d6fd1c03d525"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Dataset: Tox21_AhR_training(8169):\n",
            "====================\n",
            "Number of graphs: 8169\n",
            "Number of features: 50\n",
            "Number of classes: 2\n",
            "\n",
            "Data(edge_index=[2, 52], x=[25, 50], edge_attr=[52, 4], y=[1])\n",
            "=============================================================\n",
            "Number of nodes: 25\n",
            "Number of edges: 52\n",
            "Average node degree: 2.08\n",
            "Has isolated nodes: False\n",
            "Has self-loops: False\n",
            "Is undirected: True\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from torch_geometric.datasets import TUDataset\n",
        "\n",
        "print()\n",
        "print(f'Dataset: {dataset}:')\n",
        "print('====================')\n",
        "print(f'Number of graphs: {len(dataset)}')\n",
        "print(f'Number of features: {dataset.num_features}')\n",
        "print(f'Number of classes: {dataset.num_classes}')\n",
        "\n",
        "data = dataset[0]  # Get the first graph object.\n",
        "\n",
        "print()\n",
        "print(data)\n",
        "print('=============================================================')\n",
        "\n",
        "# Gather some statistics about the first graph.\n",
        "print(f'Number of nodes: {data.num_nodes}')\n",
        "print(f'Number of edges: {data.num_edges}')\n",
        "print(f'Average node degree: {data.num_edges / data.num_nodes:.2f}')\n",
        "print(f'Has isolated nodes: {data.has_isolated_nodes()}')\n",
        "print(f'Has self-loops: {data.has_self_loops()}')\n",
        "print(f'Is undirected: {data.is_undirected()}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "j11WiUr-PRH_",
        "outputId": "5f84abf9-05bc-4c62-95f1-5f80b603b0f0"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Number of training graphs: 6535\n",
            "Number of test graphs: 1634\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/mihirarya/Dev/bcil/gnn-drug-discovery/env/lib/python3.11/site-packages/torch_geometric/data/in_memory_dataset.py:300: UserWarning: It is not recommended to directly access the internal storage format `data` of an 'InMemoryDataset'. The data of the dataset is already cached, so any modifications to `data` will not be reflected when accessing its elements. Clearing the cache now by removing all elements in `dataset._data_list`. If you are absolutely certain what you are doing, access the internal storage via `InMemoryDataset._data` instead to suppress this warning. Alternatively, you can access stacked individual attributes of every graph via `dataset.{attr_name}`.\n",
            "  warnings.warn(msg)\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "from sklearn.model_selection import train_test_split\n",
        "from torch.utils.data import Subset\n",
        "train_idx, test_idx = train_test_split(np.arange(len(dataset)),\n",
        "                                             test_size=0.2,\n",
        "                                             shuffle=True,\n",
        "                                             stratify=dataset.data.y)\n",
        "\n",
        "train_dataset = Subset(dataset, train_idx)\n",
        "test_dataset = Subset(dataset, test_idx)\n",
        "\n",
        "\n",
        "# split = int(len(dataset) * 0.80)\n",
        "# train_dataset = dataset[:split]\n",
        "# test_dataset = dataset[split:]\n",
        "\n",
        "print(f'Number of training graphs: {len(train_dataset)}')\n",
        "print(f'Number of test graphs: {len(test_dataset)}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{0: 5775, 1: 760}"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "count = {}\n",
        "for mol in train_dataset:\n",
        "    count[mol.y[0].item()] = 1 + count.get(mol.y[0].item(), 0)\n",
        "count"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{1: 190, 0: 1444}"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "count = {}\n",
        "for mol in test_dataset:\n",
        "    count[mol.y[0].item()] = 1 + count.get(mol.y[0].item(), 0)\n",
        "count"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0gZ-l0npPIca",
        "outputId": "65b6d684-0e67-42f3-bf60-c98ec1d7338a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Step 1:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2604], x=[1271, 50], edge_attr=[2604, 4], y=[64], batch=[1271], ptr=[65])\n",
            "\n",
            "Step 2:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2432], x=[1181, 50], edge_attr=[2432, 4], y=[64], batch=[1181], ptr=[65])\n",
            "\n",
            "Step 3:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2430], x=[1184, 50], edge_attr=[2430, 4], y=[64], batch=[1184], ptr=[65])\n",
            "\n",
            "Step 4:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2342], x=[1136, 50], edge_attr=[2342, 4], y=[64], batch=[1136], ptr=[65])\n",
            "\n",
            "Step 5:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2078], x=[1028, 50], edge_attr=[2078, 4], y=[64], batch=[1028], ptr=[65])\n",
            "\n",
            "Step 6:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2612], x=[1279, 50], edge_attr=[2612, 4], y=[64], batch=[1279], ptr=[65])\n",
            "\n",
            "Step 7:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2546], x=[1225, 50], edge_attr=[2546, 4], y=[64], batch=[1225], ptr=[65])\n",
            "\n",
            "Step 8:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2164], x=[1076, 50], edge_attr=[2164, 4], y=[64], batch=[1076], ptr=[65])\n",
            "\n",
            "Step 9:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2556], x=[1237, 50], edge_attr=[2556, 4], y=[64], batch=[1237], ptr=[65])\n",
            "\n",
            "Step 10:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2658], x=[1295, 50], edge_attr=[2658, 4], y=[64], batch=[1295], ptr=[65])\n",
            "\n",
            "Step 11:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2220], x=[1087, 50], edge_attr=[2220, 4], y=[64], batch=[1087], ptr=[65])\n",
            "\n",
            "Step 12:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2304], x=[1121, 50], edge_attr=[2304, 4], y=[64], batch=[1121], ptr=[65])\n",
            "\n",
            "Step 13:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2382], x=[1170, 50], edge_attr=[2382, 4], y=[64], batch=[1170], ptr=[65])\n",
            "\n",
            "Step 14:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2478], x=[1204, 50], edge_attr=[2478, 4], y=[64], batch=[1204], ptr=[65])\n",
            "\n",
            "Step 15:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2134], x=[1052, 50], edge_attr=[2134, 4], y=[64], batch=[1052], ptr=[65])\n",
            "\n",
            "Step 16:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2474], x=[1201, 50], edge_attr=[2474, 4], y=[64], batch=[1201], ptr=[65])\n",
            "\n",
            "Step 17:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2430], x=[1195, 50], edge_attr=[2430, 4], y=[64], batch=[1195], ptr=[65])\n",
            "\n",
            "Step 18:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2284], x=[1115, 50], edge_attr=[2284, 4], y=[64], batch=[1115], ptr=[65])\n",
            "\n",
            "Step 19:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2604], x=[1270, 50], edge_attr=[2604, 4], y=[64], batch=[1270], ptr=[65])\n",
            "\n",
            "Step 20:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2402], x=[1184, 50], edge_attr=[2402, 4], y=[64], batch=[1184], ptr=[65])\n",
            "\n",
            "Step 21:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2498], x=[1206, 50], edge_attr=[2498, 4], y=[64], batch=[1206], ptr=[65])\n",
            "\n",
            "Step 22:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2430], x=[1183, 50], edge_attr=[2430, 4], y=[64], batch=[1183], ptr=[65])\n",
            "\n",
            "Step 23:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2180], x=[1072, 50], edge_attr=[2180, 4], y=[64], batch=[1072], ptr=[65])\n",
            "\n",
            "Step 24:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2594], x=[1261, 50], edge_attr=[2594, 4], y=[64], batch=[1261], ptr=[65])\n",
            "\n",
            "Step 25:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2730], x=[1331, 50], edge_attr=[2730, 4], y=[64], batch=[1331], ptr=[65])\n",
            "\n",
            "Step 26:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2370], x=[1166, 50], edge_attr=[2370, 4], y=[64], batch=[1166], ptr=[65])\n",
            "\n",
            "Step 27:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2134], x=[1054, 50], edge_attr=[2134, 4], y=[64], batch=[1054], ptr=[65])\n",
            "\n",
            "Step 28:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2474], x=[1201, 50], edge_attr=[2474, 4], y=[64], batch=[1201], ptr=[65])\n",
            "\n",
            "Step 29:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2332], x=[1146, 50], edge_attr=[2332, 4], y=[64], batch=[1146], ptr=[65])\n",
            "\n",
            "Step 30:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2062], x=[1021, 50], edge_attr=[2062, 4], y=[64], batch=[1021], ptr=[65])\n",
            "\n",
            "Step 31:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2818], x=[1366, 50], edge_attr=[2818, 4], y=[64], batch=[1366], ptr=[65])\n",
            "\n",
            "Step 32:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2362], x=[1157, 50], edge_attr=[2362, 4], y=[64], batch=[1157], ptr=[65])\n",
            "\n",
            "Step 33:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2852], x=[1389, 50], edge_attr=[2852, 4], y=[64], batch=[1389], ptr=[65])\n",
            "\n",
            "Step 34:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2182], x=[1067, 50], edge_attr=[2182, 4], y=[64], batch=[1067], ptr=[65])\n",
            "\n",
            "Step 35:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 1988], x=[988, 50], edge_attr=[1988, 4], y=[64], batch=[988], ptr=[65])\n",
            "\n",
            "Step 36:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2208], x=[1073, 50], edge_attr=[2208, 4], y=[64], batch=[1073], ptr=[65])\n",
            "\n",
            "Step 37:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2386], x=[1166, 50], edge_attr=[2386, 4], y=[64], batch=[1166], ptr=[65])\n",
            "\n",
            "Step 38:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2388], x=[1174, 50], edge_attr=[2388, 4], y=[64], batch=[1174], ptr=[65])\n",
            "\n",
            "Step 39:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2408], x=[1166, 50], edge_attr=[2408, 4], y=[64], batch=[1166], ptr=[65])\n",
            "\n",
            "Step 40:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2268], x=[1119, 50], edge_attr=[2268, 4], y=[64], batch=[1119], ptr=[65])\n",
            "\n",
            "Step 41:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2634], x=[1273, 50], edge_attr=[2634, 4], y=[64], batch=[1273], ptr=[65])\n",
            "\n",
            "Step 42:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2310], x=[1127, 50], edge_attr=[2310, 4], y=[64], batch=[1127], ptr=[65])\n",
            "\n",
            "Step 43:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2216], x=[1081, 50], edge_attr=[2216, 4], y=[64], batch=[1081], ptr=[65])\n",
            "\n",
            "Step 44:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2266], x=[1117, 50], edge_attr=[2266, 4], y=[64], batch=[1117], ptr=[65])\n",
            "\n",
            "Step 45:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2524], x=[1236, 50], edge_attr=[2524, 4], y=[64], batch=[1236], ptr=[65])\n",
            "\n",
            "Step 46:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2400], x=[1178, 50], edge_attr=[2400, 4], y=[64], batch=[1178], ptr=[65])\n",
            "\n",
            "Step 47:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2386], x=[1179, 50], edge_attr=[2386, 4], y=[64], batch=[1179], ptr=[65])\n",
            "\n",
            "Step 48:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2560], x=[1231, 50], edge_attr=[2560, 4], y=[64], batch=[1231], ptr=[65])\n",
            "\n",
            "Step 49:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2412], x=[1180, 50], edge_attr=[2412, 4], y=[64], batch=[1180], ptr=[65])\n",
            "\n",
            "Step 50:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2110], x=[1033, 50], edge_attr=[2110, 4], y=[64], batch=[1033], ptr=[65])\n",
            "\n",
            "Step 51:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2342], x=[1142, 50], edge_attr=[2342, 4], y=[64], batch=[1142], ptr=[65])\n",
            "\n",
            "Step 52:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2304], x=[1127, 50], edge_attr=[2304, 4], y=[64], batch=[1127], ptr=[65])\n",
            "\n",
            "Step 53:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2172], x=[1053, 50], edge_attr=[2172, 4], y=[64], batch=[1053], ptr=[65])\n",
            "\n",
            "Step 54:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2674], x=[1292, 50], edge_attr=[2674, 4], y=[64], batch=[1292], ptr=[65])\n",
            "\n",
            "Step 55:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2494], x=[1217, 50], edge_attr=[2494, 4], y=[64], batch=[1217], ptr=[65])\n",
            "\n",
            "Step 56:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2612], x=[1267, 50], edge_attr=[2612, 4], y=[64], batch=[1267], ptr=[65])\n",
            "\n",
            "Step 57:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2058], x=[1004, 50], edge_attr=[2058, 4], y=[64], batch=[1004], ptr=[65])\n",
            "\n",
            "Step 58:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2316], x=[1127, 50], edge_attr=[2316, 4], y=[64], batch=[1127], ptr=[65])\n",
            "\n",
            "Step 59:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2340], x=[1136, 50], edge_attr=[2340, 4], y=[64], batch=[1136], ptr=[65])\n",
            "\n",
            "Step 60:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2184], x=[1093, 50], edge_attr=[2184, 4], y=[64], batch=[1093], ptr=[65])\n",
            "\n",
            "Step 61:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2270], x=[1111, 50], edge_attr=[2270, 4], y=[64], batch=[1111], ptr=[65])\n",
            "\n",
            "Step 62:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2246], x=[1105, 50], edge_attr=[2246, 4], y=[64], batch=[1105], ptr=[65])\n",
            "\n",
            "Step 63:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2016], x=[996, 50], edge_attr=[2016, 4], y=[64], batch=[996], ptr=[65])\n",
            "\n",
            "Step 64:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 1864], x=[938, 50], edge_attr=[1864, 4], y=[64], batch=[938], ptr=[65])\n",
            "\n",
            "Step 65:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2874], x=[1382, 50], edge_attr=[2874, 4], y=[64], batch=[1382], ptr=[65])\n",
            "\n",
            "Step 66:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2110], x=[1037, 50], edge_attr=[2110, 4], y=[64], batch=[1037], ptr=[65])\n",
            "\n",
            "Step 67:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2418], x=[1196, 50], edge_attr=[2418, 4], y=[64], batch=[1196], ptr=[65])\n",
            "\n",
            "Step 68:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2834], x=[1385, 50], edge_attr=[2834, 4], y=[64], batch=[1385], ptr=[65])\n",
            "\n",
            "Step 69:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2270], x=[1104, 50], edge_attr=[2270, 4], y=[64], batch=[1104], ptr=[65])\n",
            "\n",
            "Step 70:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2186], x=[1066, 50], edge_attr=[2186, 4], y=[64], batch=[1066], ptr=[65])\n",
            "\n",
            "Step 71:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 1926], x=[968, 50], edge_attr=[1926, 4], y=[64], batch=[968], ptr=[65])\n",
            "\n",
            "Step 72:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2048], x=[1020, 50], edge_attr=[2048, 4], y=[64], batch=[1020], ptr=[65])\n",
            "\n",
            "Step 73:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2308], x=[1137, 50], edge_attr=[2308, 4], y=[64], batch=[1137], ptr=[65])\n",
            "\n",
            "Step 74:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2460], x=[1199, 50], edge_attr=[2460, 4], y=[64], batch=[1199], ptr=[65])\n",
            "\n",
            "Step 75:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2284], x=[1117, 50], edge_attr=[2284, 4], y=[64], batch=[1117], ptr=[65])\n",
            "\n",
            "Step 76:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2504], x=[1223, 50], edge_attr=[2504, 4], y=[64], batch=[1223], ptr=[65])\n",
            "\n",
            "Step 77:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2216], x=[1093, 50], edge_attr=[2216, 4], y=[64], batch=[1093], ptr=[65])\n",
            "\n",
            "Step 78:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2236], x=[1091, 50], edge_attr=[2236, 4], y=[64], batch=[1091], ptr=[65])\n",
            "\n",
            "Step 79:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2038], x=[1010, 50], edge_attr=[2038, 4], y=[64], batch=[1010], ptr=[65])\n",
            "\n",
            "Step 80:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2186], x=[1092, 50], edge_attr=[2186, 4], y=[64], batch=[1092], ptr=[65])\n",
            "\n",
            "Step 81:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2306], x=[1126, 50], edge_attr=[2306, 4], y=[64], batch=[1126], ptr=[65])\n",
            "\n",
            "Step 82:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2380], x=[1163, 50], edge_attr=[2380, 4], y=[64], batch=[1163], ptr=[65])\n",
            "\n",
            "Step 83:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2324], x=[1136, 50], edge_attr=[2324, 4], y=[64], batch=[1136], ptr=[65])\n",
            "\n",
            "Step 84:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2442], x=[1197, 50], edge_attr=[2442, 4], y=[64], batch=[1197], ptr=[65])\n",
            "\n",
            "Step 85:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2752], x=[1336, 50], edge_attr=[2752, 4], y=[64], batch=[1336], ptr=[65])\n",
            "\n",
            "Step 86:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2400], x=[1181, 50], edge_attr=[2400, 4], y=[64], batch=[1181], ptr=[65])\n",
            "\n",
            "Step 87:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2058], x=[1020, 50], edge_attr=[2058, 4], y=[64], batch=[1020], ptr=[65])\n",
            "\n",
            "Step 88:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2560], x=[1241, 50], edge_attr=[2560, 4], y=[64], batch=[1241], ptr=[65])\n",
            "\n",
            "Step 89:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2572], x=[1230, 50], edge_attr=[2572, 4], y=[64], batch=[1230], ptr=[65])\n",
            "\n",
            "Step 90:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2362], x=[1153, 50], edge_attr=[2362, 4], y=[64], batch=[1153], ptr=[65])\n",
            "\n",
            "Step 91:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2270], x=[1111, 50], edge_attr=[2270, 4], y=[64], batch=[1111], ptr=[65])\n",
            "\n",
            "Step 92:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2450], x=[1198, 50], edge_attr=[2450, 4], y=[64], batch=[1198], ptr=[65])\n",
            "\n",
            "Step 93:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2360], x=[1151, 50], edge_attr=[2360, 4], y=[64], batch=[1151], ptr=[65])\n",
            "\n",
            "Step 94:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2510], x=[1203, 50], edge_attr=[2510, 4], y=[64], batch=[1203], ptr=[65])\n",
            "\n",
            "Step 95:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2434], x=[1198, 50], edge_attr=[2434, 4], y=[64], batch=[1198], ptr=[65])\n",
            "\n",
            "Step 96:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2388], x=[1186, 50], edge_attr=[2388, 4], y=[64], batch=[1186], ptr=[65])\n",
            "\n",
            "Step 97:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2522], x=[1232, 50], edge_attr=[2522, 4], y=[64], batch=[1232], ptr=[65])\n",
            "\n",
            "Step 98:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2424], x=[1190, 50], edge_attr=[2424, 4], y=[64], batch=[1190], ptr=[65])\n",
            "\n",
            "Step 99:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2578], x=[1254, 50], edge_attr=[2578, 4], y=[64], batch=[1254], ptr=[65])\n",
            "\n",
            "Step 100:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2708], x=[1299, 50], edge_attr=[2708, 4], y=[64], batch=[1299], ptr=[65])\n",
            "\n",
            "Step 101:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2394], x=[1178, 50], edge_attr=[2394, 4], y=[64], batch=[1178], ptr=[65])\n",
            "\n",
            "Step 102:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2450], x=[1186, 50], edge_attr=[2450, 4], y=[64], batch=[1186], ptr=[65])\n",
            "\n",
            "Step 103:\n",
            "=======\n",
            "Number of graphs in the current batch: 7\n",
            "DataBatch(edge_index=[2, 220], x=[109, 50], edge_attr=[220, 4], y=[7], batch=[109], ptr=[8])\n",
            "\n"
          ]
        }
      ],
      "source": [
        "from torch_geometric.loader import DataLoader\n",
        "\n",
        "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
        "test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)\n",
        "\n",
        "for step, data in enumerate(train_loader):\n",
        "    print(f'Step {step + 1}:')\n",
        "    print('=======')\n",
        "    print(f'Number of graphs in the current batch: {data.num_graphs}')\n",
        "    print(data)\n",
        "    print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Step 1:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2188], x=[1072, 50], edge_attr=[2188, 4], y=[64], batch=[1072], ptr=[65])\n",
            "\n",
            "Step 2:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2366], x=[1153, 50], edge_attr=[2366, 4], y=[64], batch=[1153], ptr=[65])\n",
            "\n",
            "Step 3:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2366], x=[1165, 50], edge_attr=[2366, 4], y=[64], batch=[1165], ptr=[65])\n",
            "\n",
            "Step 4:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2220], x=[1083, 50], edge_attr=[2220, 4], y=[64], batch=[1083], ptr=[65])\n",
            "\n",
            "Step 5:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2370], x=[1154, 50], edge_attr=[2370, 4], y=[64], batch=[1154], ptr=[65])\n",
            "\n",
            "Step 6:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2498], x=[1222, 50], edge_attr=[2498, 4], y=[64], batch=[1222], ptr=[65])\n",
            "\n",
            "Step 7:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2494], x=[1222, 50], edge_attr=[2494, 4], y=[64], batch=[1222], ptr=[65])\n",
            "\n",
            "Step 8:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2206], x=[1088, 50], edge_attr=[2206, 4], y=[64], batch=[1088], ptr=[65])\n",
            "\n",
            "Step 9:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2128], x=[1069, 50], edge_attr=[2128, 4], y=[64], batch=[1069], ptr=[65])\n",
            "\n",
            "Step 10:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2146], x=[1048, 50], edge_attr=[2146, 4], y=[64], batch=[1048], ptr=[65])\n",
            "\n",
            "Step 11:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2364], x=[1154, 50], edge_attr=[2364, 4], y=[64], batch=[1154], ptr=[65])\n",
            "\n",
            "Step 12:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2238], x=[1085, 50], edge_attr=[2238, 4], y=[64], batch=[1085], ptr=[65])\n",
            "\n",
            "Step 13:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2382], x=[1172, 50], edge_attr=[2382, 4], y=[64], batch=[1172], ptr=[65])\n",
            "\n",
            "Step 14:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2448], x=[1182, 50], edge_attr=[2448, 4], y=[64], batch=[1182], ptr=[65])\n",
            "\n",
            "Step 15:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2634], x=[1280, 50], edge_attr=[2634, 4], y=[64], batch=[1280], ptr=[65])\n",
            "\n",
            "Step 16:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2192], x=[1077, 50], edge_attr=[2192, 4], y=[64], batch=[1077], ptr=[65])\n",
            "\n",
            "Step 17:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2436], x=[1192, 50], edge_attr=[2436, 4], y=[64], batch=[1192], ptr=[65])\n",
            "\n",
            "Step 18:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2108], x=[1036, 50], edge_attr=[2108, 4], y=[64], batch=[1036], ptr=[65])\n",
            "\n",
            "Step 19:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2396], x=[1166, 50], edge_attr=[2396, 4], y=[64], batch=[1166], ptr=[65])\n",
            "\n",
            "Step 20:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2210], x=[1091, 50], edge_attr=[2210, 4], y=[64], batch=[1091], ptr=[65])\n",
            "\n",
            "Step 21:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2408], x=[1163, 50], edge_attr=[2408, 4], y=[64], batch=[1163], ptr=[65])\n",
            "\n",
            "Step 22:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2090], x=[1040, 50], edge_attr=[2090, 4], y=[64], batch=[1040], ptr=[65])\n",
            "\n",
            "Step 23:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2246], x=[1106, 50], edge_attr=[2246, 4], y=[64], batch=[1106], ptr=[65])\n",
            "\n",
            "Step 24:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2622], x=[1280, 50], edge_attr=[2622, 4], y=[64], batch=[1280], ptr=[65])\n",
            "\n",
            "Step 25:\n",
            "=======\n",
            "Number of graphs in the current batch: 64\n",
            "DataBatch(edge_index=[2, 2984], x=[1445, 50], edge_attr=[2984, 4], y=[64], batch=[1445], ptr=[65])\n",
            "\n",
            "Step 26:\n",
            "=======\n",
            "Number of graphs in the current batch: 34\n",
            "DataBatch(edge_index=[2, 1190], x=[576, 50], edge_attr=[1190, 4], y=[34], batch=[576], ptr=[35])\n",
            "\n"
          ]
        }
      ],
      "source": [
        "for step, data in enumerate(test_loader):\n",
        "    print(f'Step {step + 1}:')\n",
        "    print('=======')\n",
        "    print(f'Number of graphs in the current batch: {data.num_graphs}')\n",
        "    print(data)\n",
        "    print()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CN3sRVuaQ88l",
        "outputId": "59c998ba-b387-4413-d7d3-2101a27495bd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "GCN(\n",
            "  (conv1): GCNConv(50, 64)\n",
            "  (conv2): GCNConv(64, 64)\n",
            "  (conv3): GCNConv(64, 64)\n",
            "  (lin): Linear(in_features=64, out_features=2, bias=True)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "from torch.nn import Linear\n",
        "import torch.nn.functional as F\n",
        "from torch_geometric.nn import GCNConv\n",
        "from torch_geometric.nn import global_mean_pool, global_max_pool, global_add_pool\n",
        "\n",
        "\n",
        "class GCN(torch.nn.Module):\n",
        "    def __init__(self, hidden_channels):\n",
        "        super(GCN, self).__init__()\n",
        "        self.conv1 = GCNConv(dataset.num_node_features, hidden_channels)\n",
        "        self.conv2 = GCNConv(hidden_channels, hidden_channels)\n",
        "        self.conv3 = GCNConv(hidden_channels, hidden_channels)\n",
        "        self.lin = Linear(hidden_channels, dataset.num_classes)\n",
        "\n",
        "    def forward(self, x, edge_index, batch):\n",
        "        # 1. Obtain node embeddings\n",
        "        x = self.conv1(x, edge_index)\n",
        "        x = x.relu()\n",
        "        x = self.conv2(x, edge_index)\n",
        "        x = x.relu()\n",
        "        x = self.conv3(x, edge_index)\n",
        "\n",
        "        # 2. Readout layer\n",
        "        x = global_add_pool(x, batch)  # [batch_size, hidden_channels]\n",
        "\n",
        "        # 3. Apply a final classifier\n",
        "        x = F.dropout(x, p=0.5, training=self.training)\n",
        "        x = self.lin(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "model = GCN(hidden_channels=64)\n",
        "print(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 300
        },
        "id": "HvhgQoO8Svw4",
        "outputId": "4d17af35-85c8-4bcd-8e16-c514fd2ff714"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch: 001, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.7172\n",
            "Epoch: 002, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8069\n",
            "Epoch: 003, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8276\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/Users/mihirarya/Dev/bcil/gnn-drug-discovery/env/lib/python3.11/site-packages/torchmetrics/utilities/prints.py:43: UserWarning: No positive samples in targets, true positive value should be meaningless. Returning zero tensor in true positive score\n",
            "  warnings.warn(*args, **kwargs)  # noqa: B028\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Epoch: 004, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8690\n",
            "Epoch: 005, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8966\n",
            "Epoch: 006, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.9172\n",
            "Epoch: 007, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8690\n",
            "Epoch: 008, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8759\n",
            "Epoch: 009, Train Acc: 0.8828, Test Acc: 0.8843, Test AUC: 0.8000\n",
            "Epoch: 010, Train Acc: 0.8799, Test Acc: 0.8807, Test AUC: 0.8000\n",
            "Epoch: 011, Train Acc: 0.8817, Test Acc: 0.8825, Test AUC: 0.8276\n",
            "Epoch: 012, Train Acc: 0.8831, Test Acc: 0.8825, Test AUC: 0.8069\n",
            "Epoch: 013, Train Acc: 0.8842, Test Acc: 0.8831, Test AUC: 0.8207\n",
            "Epoch: 014, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.6552\n",
            "Epoch: 015, Train Acc: 0.8840, Test Acc: 0.8813, Test AUC: 0.8000\n",
            "Epoch: 016, Train Acc: 0.8834, Test Acc: 0.8813, Test AUC: 0.7724\n",
            "Epoch: 017, Train Acc: 0.8840, Test Acc: 0.8825, Test AUC: 0.8414\n",
            "Epoch: 018, Train Acc: 0.8846, Test Acc: 0.8831, Test AUC: 0.8621\n",
            "Epoch: 019, Train Acc: 0.8840, Test Acc: 0.8825, Test AUC: 0.8483\n",
            "Epoch: 020, Train Acc: 0.8848, Test Acc: 0.8819, Test AUC: 0.7034\n",
            "Epoch: 021, Train Acc: 0.8854, Test Acc: 0.8868, Test AUC: 0.7862\n",
            "Epoch: 022, Train Acc: 0.8842, Test Acc: 0.8813, Test AUC: 0.8207\n",
            "Epoch: 023, Train Acc: 0.8832, Test Acc: 0.8813, Test AUC: 0.8828\n",
            "Epoch: 024, Train Acc: 0.8865, Test Acc: 0.8837, Test AUC: 0.8138\n",
            "Epoch: 025, Train Acc: 0.8846, Test Acc: 0.8837, Test AUC: 0.8621\n",
            "Epoch: 026, Train Acc: 0.8842, Test Acc: 0.8819, Test AUC: 0.8483\n",
            "Epoch: 027, Train Acc: 0.8877, Test Acc: 0.8825, Test AUC: 0.8207\n",
            "Epoch: 028, Train Acc: 0.8854, Test Acc: 0.8868, Test AUC: 0.8414\n",
            "Epoch: 029, Train Acc: 0.8837, Test Acc: 0.8807, Test AUC: 0.9310\n",
            "Epoch: 030, Train Acc: 0.8865, Test Acc: 0.8843, Test AUC: 0.9034\n",
            "Epoch: 031, Train Acc: 0.8836, Test Acc: 0.8831, Test AUC: 0.9034\n",
            "Epoch: 032, Train Acc: 0.8855, Test Acc: 0.8813, Test AUC: 0.8759\n",
            "Epoch: 033, Train Acc: 0.8845, Test Acc: 0.8782, Test AUC: 0.9241\n",
            "Epoch: 034, Train Acc: 0.8851, Test Acc: 0.8825, Test AUC: 0.9241\n",
            "Epoch: 035, Train Acc: 0.8839, Test Acc: 0.8837, Test AUC: 0.8552\n",
            "Epoch: 036, Train Acc: 0.8813, Test Acc: 0.8794, Test AUC: 0.9379\n",
            "Epoch: 037, Train Acc: 0.8816, Test Acc: 0.8678, Test AUC: 0.9310\n",
            "Epoch: 038, Train Acc: 0.8862, Test Acc: 0.8837, Test AUC: 0.8690\n",
            "Epoch: 039, Train Acc: 0.8875, Test Acc: 0.8813, Test AUC: 0.9034\n",
            "Epoch: 040, Train Acc: 0.8894, Test Acc: 0.8849, Test AUC: 0.8759\n",
            "Epoch: 041, Train Acc: 0.8888, Test Acc: 0.8837, Test AUC: 0.9241\n",
            "Epoch: 042, Train Acc: 0.8813, Test Acc: 0.8739, Test AUC: 0.9310\n",
            "Epoch: 043, Train Acc: 0.8865, Test Acc: 0.8880, Test AUC: 0.8345\n",
            "Epoch: 044, Train Acc: 0.8929, Test Acc: 0.8794, Test AUC: 0.8966\n",
            "Epoch: 045, Train Acc: 0.8842, Test Acc: 0.8862, Test AUC: 0.9172\n",
            "Epoch: 046, Train Acc: 0.8862, Test Acc: 0.8874, Test AUC: 0.9310\n",
            "Epoch: 047, Train Acc: 0.8869, Test Acc: 0.8831, Test AUC: 0.9034\n",
            "Epoch: 048, Train Acc: 0.8892, Test Acc: 0.8868, Test AUC: 0.9172\n",
            "Epoch: 049, Train Acc: 0.8880, Test Acc: 0.8862, Test AUC: 0.9241\n"
          ]
        }
      ],
      "source": [
        "model = GCN(hidden_channels=64)\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "auroc = AUROC(task=\"binary\")\n",
        "\n",
        "def train():\n",
        "    model.train()\n",
        "\n",
        "    for data in train_loader:  # Iterate in batches over the training dataset.\n",
        "         out = model(data.x, data.edge_index, data.batch)  # Perform a single forward pass.\n",
        "         loss = criterion(out, data.y)  # Compute the loss.\n",
        "         loss.backward()  # Derive gradients.\n",
        "         optimizer.step()  # Update parameters based on gradients.\n",
        "         optimizer.zero_grad()  # Clear gradients.\n",
        "\n",
        "def test(loader):\n",
        "     model.eval()\n",
        "\n",
        "     correct = 0\n",
        "     for data in loader:  # Iterate in batches over the training/test dataset.\n",
        "         out = model(data.x, data.edge_index, data.batch)\n",
        "         pred = out.argmax(dim=1)  # Use the class with highest probability.\n",
        "         correct += int((pred == data.y).sum())  # Check against ground-truth labels.\n",
        "     return (correct / len(loader.dataset), auroc(out.softmax(dim=1)[:, 1], data.y))  # Derive ratio of correct predictions.\n",
        "\n",
        "\n",
        "for epoch in range(1, 50):\n",
        "    train()\n",
        "    train_acc, train_auc = test(train_loader)\n",
        "    test_acc, test_auc = test(test_loader)\n",
        "    print(f'Epoch: {epoch:03d}, Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}, Test AUC: {test_auc:.4f}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "GAT(\n",
            "  (conv1): GATConv(50, 64, heads=1)\n",
            "  (conv2): GATConv(64, 64, heads=1)\n",
            "  (conv3): GATConv(64, 64, heads=1)\n",
            "  (lin): Linear(in_features=64, out_features=2, bias=True)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "from torch.nn import Linear\n",
        "import torch.nn.functional as F\n",
        "from torch_geometric.nn import GATConv\n",
        "from torch_geometric.nn import global_mean_pool, global_max_pool, global_add_pool\n",
        "\n",
        "\n",
        "class GAT(torch.nn.Module):\n",
        "    def __init__(self, hidden_channels):\n",
        "        super(GAT, self).__init__()\n",
        "        self.conv1 = GATConv(dataset.num_node_features, hidden_channels)\n",
        "        self.conv2 = GATConv(hidden_channels, hidden_channels)\n",
        "        self.conv3 = GATConv(hidden_channels, hidden_channels)\n",
        "        self.lin = Linear(hidden_channels, dataset.num_classes)\n",
        "\n",
        "    def forward(self, x, edge_index, batch):\n",
        "        # 1. Obtain node embeddings\n",
        "        x = self.conv1(x, edge_index)\n",
        "        x = x.relu()\n",
        "        x = self.conv2(x, edge_index)\n",
        "        x = x.relu()\n",
        "        x = self.conv3(x, edge_index)\n",
        "\n",
        "        # 2. Readout layer\n",
        "        x = global_add_pool(x, batch)  # [batch_size, hidden_channels]\n",
        "\n",
        "        # 3. Apply a final classifier\n",
        "        x = F.dropout(x, p=0.5, training=self.training)\n",
        "        x = self.lin(x)\n",
        "\n",
        "        return x\n",
        "\n",
        "model = GAT(hidden_channels=64)\n",
        "print(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "xs5D713Ia_Sv"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "GAT(\n",
            "  (conv1): GATConv(50, 64, heads=1)\n",
            "  (conv2): GATConv(64, 64, heads=1)\n",
            "  (conv3): GATConv(64, 64, heads=1)\n",
            "  (lin): Linear(in_features=64, out_features=2, bias=True)\n",
            ")\n",
            "Epoch: 001, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8138\n",
            "Epoch: 002, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8414\n",
            "Epoch: 003, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8345\n",
            "Epoch: 004, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8828\n",
            "Epoch: 005, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.9103\n",
            "Epoch: 006, Train Acc: 0.8836, Test Acc: 0.8831, Test AUC: 0.8759\n",
            "Epoch: 007, Train Acc: 0.8849, Test Acc: 0.8856, Test AUC: 0.9034\n",
            "Epoch: 008, Train Acc: 0.8837, Test Acc: 0.8837, Test AUC: 0.8897\n",
            "Epoch: 009, Train Acc: 0.8854, Test Acc: 0.8874, Test AUC: 0.6483\n",
            "Epoch: 010, Train Acc: 0.8832, Test Acc: 0.8825, Test AUC: 0.8414\n",
            "Epoch: 011, Train Acc: 0.8846, Test Acc: 0.8868, Test AUC: 0.8138\n",
            "Epoch: 012, Train Acc: 0.8860, Test Acc: 0.8856, Test AUC: 0.9103\n",
            "Epoch: 013, Train Acc: 0.8858, Test Acc: 0.8843, Test AUC: 0.7517\n",
            "Epoch: 014, Train Acc: 0.8862, Test Acc: 0.8856, Test AUC: 0.8759\n",
            "Epoch: 015, Train Acc: 0.8878, Test Acc: 0.8862, Test AUC: 0.8690\n",
            "Epoch: 016, Train Acc: 0.8880, Test Acc: 0.8886, Test AUC: 0.9172\n",
            "Epoch: 017, Train Acc: 0.8901, Test Acc: 0.8874, Test AUC: 0.8897\n",
            "Epoch: 018, Train Acc: 0.8883, Test Acc: 0.8880, Test AUC: 0.9448\n",
            "Epoch: 019, Train Acc: 0.8880, Test Acc: 0.8880, Test AUC: 0.9448\n",
            "Epoch: 020, Train Acc: 0.8880, Test Acc: 0.8868, Test AUC: 0.8345\n",
            "Epoch: 021, Train Acc: 0.8872, Test Acc: 0.8849, Test AUC: 0.9172\n",
            "Epoch: 022, Train Acc: 0.8900, Test Acc: 0.8892, Test AUC: 0.8345\n",
            "Epoch: 023, Train Acc: 0.8904, Test Acc: 0.8874, Test AUC: 0.8966\n",
            "Epoch: 024, Train Acc: 0.8906, Test Acc: 0.8886, Test AUC: 0.9241\n",
            "Epoch: 025, Train Acc: 0.8898, Test Acc: 0.8880, Test AUC: 0.9103\n",
            "Epoch: 026, Train Acc: 0.8936, Test Acc: 0.8880, Test AUC: 0.8897\n",
            "Epoch: 027, Train Acc: 0.8891, Test Acc: 0.8862, Test AUC: 0.9172\n",
            "Epoch: 028, Train Acc: 0.8892, Test Acc: 0.8843, Test AUC: 0.9379\n",
            "Epoch: 029, Train Acc: 0.8921, Test Acc: 0.8886, Test AUC: 0.9517\n",
            "Epoch: 030, Train Acc: 0.8923, Test Acc: 0.8837, Test AUC: 0.9517\n",
            "Epoch: 031, Train Acc: 0.8877, Test Acc: 0.8898, Test AUC: 0.9310\n",
            "Epoch: 032, Train Acc: 0.8894, Test Acc: 0.8880, Test AUC: 0.9241\n",
            "Epoch: 033, Train Acc: 0.8865, Test Acc: 0.8849, Test AUC: 0.9379\n",
            "Epoch: 034, Train Acc: 0.8880, Test Acc: 0.8874, Test AUC: 0.9310\n",
            "Epoch: 035, Train Acc: 0.8889, Test Acc: 0.8856, Test AUC: 0.9310\n",
            "Epoch: 036, Train Acc: 0.8910, Test Acc: 0.8862, Test AUC: 0.9103\n",
            "Epoch: 037, Train Acc: 0.8912, Test Acc: 0.8923, Test AUC: 0.8897\n",
            "Epoch: 038, Train Acc: 0.8889, Test Acc: 0.8898, Test AUC: 0.8828\n",
            "Epoch: 039, Train Acc: 0.8903, Test Acc: 0.8874, Test AUC: 0.9241\n",
            "Epoch: 040, Train Acc: 0.8883, Test Acc: 0.8898, Test AUC: 0.9517\n",
            "Epoch: 041, Train Acc: 0.8894, Test Acc: 0.8892, Test AUC: 0.9448\n",
            "Epoch: 042, Train Acc: 0.8921, Test Acc: 0.8849, Test AUC: 0.9448\n",
            "Epoch: 043, Train Acc: 0.8872, Test Acc: 0.8849, Test AUC: 0.9448\n",
            "Epoch: 044, Train Acc: 0.8920, Test Acc: 0.8905, Test AUC: 0.9103\n",
            "Epoch: 045, Train Acc: 0.8866, Test Acc: 0.8898, Test AUC: 0.9448\n",
            "Epoch: 046, Train Acc: 0.8884, Test Acc: 0.8874, Test AUC: 0.9517\n",
            "Epoch: 047, Train Acc: 0.8918, Test Acc: 0.8917, Test AUC: 0.9172\n",
            "Epoch: 048, Train Acc: 0.8918, Test Acc: 0.8874, Test AUC: 0.9103\n",
            "Epoch: 049, Train Acc: 0.8901, Test Acc: 0.8892, Test AUC: 0.9379\n"
          ]
        }
      ],
      "source": [
        "train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)\n",
        "test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)\n",
        "model = GAT(hidden_channels=64)\n",
        "print(model)\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n",
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "auroc = AUROC(task=\"binary\")\n",
        "\n",
        "def train():\n",
        "    model.train()\n",
        "\n",
        "    for data in train_loader:  # Iterate in batches over the training dataset.\n",
        "         out = model(data.x, data.edge_index, data.batch)  # Perform a single forward pass.\n",
        "         loss = criterion(out, data.y)  # Compute the loss.\n",
        "         loss.backward()  # Derive gradients.\n",
        "         optimizer.step()  # Update parameters based on gradients.\n",
        "         optimizer.zero_grad()  # Clear gradients.\n",
        "\n",
        "def test(loader):\n",
        "     model.eval()\n",
        "\n",
        "     correct = 0\n",
        "     for data in loader:  # Iterate in batches over the training/test dataset.\n",
        "         out = model(data.x, data.edge_index, data.batch)\n",
        "         pred = out.argmax(dim=1)  # Use the class with highest probability.\n",
        "         correct += int((pred == data.y).sum())  # Check against ground-truth labels.\n",
        "     return (correct / len(loader.dataset), auroc(out.softmax(dim=1)[:, 1], data.y))  # Derive ratio of correct predictions.\n",
        "\n",
        "for epoch in range(1, 50):\n",
        "    train()\n",
        "    train_acc, train_auc = test(train_loader)\n",
        "    test_acc, test_auc = test(test_loader)\n",
        "    print(f'Epoch: {epoch:03d}, Train Acc: {train_acc:.4f}, Test Acc: {test_acc:.4f}, Test AUC: {test_auc:.4f}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.5"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}