notebooks/graph_regress.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graph-level regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "import sys\n",
    "sys.path.append('..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch_geometric import datasets\n",
    "from torch_geometric.loader import DataLoader\n",
    "\n",
    "from gnn_utils import (\n",
    "    GCNModel,\n",
    "    train_graph_level,\n",
    "    plot_training_curves\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set random seed\n",
    "_ = torch.manual_seed(12345)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load dataset with node features and graph targets\n",
    "data_set = datasets.QM9(root='../data')\n",
    "\n",
    "# create dataloaders\n",
    "train_loader = DataLoader(\n",
    "    data_set[:100000],\n",
    "    batch_size=32,\n",
    "    shuffle=True\n",
    ")\n",
    "\n",
    "val_loader = DataLoader(\n",
    "    data_set[100000:],\n",
    "    batch_size=32,\n",
    "    shuffle=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print summaries\n",
    "print(f'Number of graphs: {len(data_set)}')\n",
    "print(f'Number of node features: {data_set.num_features}')\n",
    "print(f'Number of graph targets: {data_set.num_classes}')\n",
    "\n",
    "print(f'\\nTensor shapes (single graph):\\n{data_set[0]}')\n",
    "print(f'Features shape: {data_set[0].x.shape}') # (num_nodes, num_features)\n",
    "print(f'Targets shape: {data_set[0].y.shape}') # (1, num_targets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get batch (as a single graph)\n",
    "batch = next(iter(val_loader)) # batching works by assembling a larger graph with isolated subgraphs\n",
    "\n",
    "# print summaries\n",
    "print(f'Tensor shapes (batch):\\n{batch}')\n",
    "print(f'Number of graphs: {batch.num_graphs}')\n",
    "print(f'Number of nodes: {batch.num_nodes}')\n",
    "print(f'Number of edges: {batch.num_edges}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print indices locating nodes within the batch\n",
    "print(f'Nodes-in-batch indices shape: {batch.batch.shape}') # (num_nodes,)\n",
    "# print(f'Nodes-in-batch indices:\\n{batch.batch}') # values in {0,...,batch_size-1}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create model\n",
    "model = GCNModel(\n",
    "    num_channels=[data_set.num_features, 128, 16],\n",
    "    num_features=data_set.num_classes,\n",
    "    graph_level=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check output shape\n",
    "y = model(batch.x, batch.edge_index, batch.batch)\n",
    "\n",
    "print(f'Node features shape: {batch.x.shape}') # (num_nodes, num_features)\n",
    "print(f'Prediction shape: {y.shape}') # (batch_size, num_targets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# determine device\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create loss function\n",
    "criterion = nn.MSELoss(reduction='mean')\n",
    "\n",
    "# initialize optimizer\n",
    "optimizer = torch.optim.Adam(\n",
    "    model.parameters(),\n",
    "    lr=0.01,\n",
    "    weight_decay=0.01\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run training\n",
    "history = train_graph_level(\n",
    "    model=model,\n",
    "    criterion=criterion,\n",
    "    optimizer=optimizer,\n",
    "    num_epochs=50,\n",
    "    train_loader=train_loader,\n",
    "    val_loader=val_loader,\n",
    "    log_every=1\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot learning curves\n",
    "fig, axes = plot_training_curves(history, figsize=(5, 4))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}