wehub-resource-sync
411 lines
18 KiB
Plaintext
411 lines
18 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "Jv-tHPvR-JKa"
|
|
},
|
|
"source": [
|
|
"# Graph Transformer in a Nutshell\n",
|
|
"\n",
|
|
"The **Transformer** [(Vaswani et al. 2017)](https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html) has been proven an effective learning architecture in natural language processing and computer vision.\n",
|
|
"Recently, researchers turns to explore the application of transformer in graph learning. They have achieved inital success on many practical tasks, e.g., graph property prediction.\n",
|
|
"[Dwivedi et al. (2020)](https://arxiv.org/abs/2012.09699) firstly generalize the transformer neural architecture to graph-structured data. Here, we present how to build such a graph transformer with DGL's sparse matrix APIs.\n",
|
|
"\n",
|
|
"[](https://colab.research.google.com/github/dmlc/dgl/blob/master/notebooks/sparse/graph_transformer.ipynb) [](https://github.com/dmlc/dgl/blob/master/notebooks/sparse/graph_transformer.ipynb)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"source": [
|
|
"# Install required packages.\n",
|
|
"import os\n",
|
|
"import torch\n",
|
|
"os.environ['TORCH'] = torch.__version__\n",
|
|
"os.environ['DGLBACKEND'] = \"pytorch\"\n",
|
|
"\n",
|
|
"# Uncomment below to install required packages. If the CUDA version is not 11.8,\n",
|
|
"# check the https://www.dgl.ai/pages/start.html to find the supported CUDA\n",
|
|
"# version and corresponding command to install DGL.\n",
|
|
"#!pip install dgl -f https://data.dgl.ai/wheels/cu118/repo.html > /dev/null\n",
|
|
"#!pip install ogb >/dev/null\n",
|
|
"\n",
|
|
"try:\n",
|
|
" import dgl\n",
|
|
" installed = True\n",
|
|
"except ImportError:\n",
|
|
" installed = False\n",
|
|
"print(\"DGL installed!\" if installed else \"Failed to install DGL!\")"
|
|
],
|
|
"metadata": {
|
|
"id": "8wIJZQqODy-7"
|
|
},
|
|
"execution_count": null,
|
|
"outputs": []
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "nOpFdtLI-JKb"
|
|
},
|
|
"source": [
|
|
"## Sparse Multi-head Attention\n",
|
|
"\n",
|
|
"Recall the all-pairs scaled-dot-product attention mechanism in vanillar Transformer:\n",
|
|
"\n",
|
|
"$$\\text{Attn}=\\text{softmax}(\\dfrac{QK^T} {\\sqrt{d}})V,$$\n",
|
|
"\n",
|
|
"The graph transformer (GT) model employs a Sparse Multi-head Attention block:\n",
|
|
"\n",
|
|
"$$\\text{SparseAttn}(Q, K, V, A) = \\text{softmax}(\\frac{(QK^T) \\circ A}{\\sqrt{d}})V,$$\n",
|
|
"\n",
|
|
"where $Q, K, V ∈\\mathbb{R}^{N\\times d}$ are query feature, key feature, and value feature, respectively. $A\\in[0,1]^{N\\times N}$ is the adjacency matrix of the input graph. $(QK^T)\\circ A$ means that the multiplication of query matrix and key matrix is followed by a Hadamard product (or element-wise multiplication) with the sparse adjacency matrix as illustrated in the figure below:\n",
|
|
"\n",
|
|
"<img src=\"https://drive.google.com/uc?id=1OgMAewLR3Z1vz5y4J8aPRSeaU3g8iQfX\" width=\"500\">\n",
|
|
"\n",
|
|
"Essentially, only the attention scores between connected nodes are computed according to the sparsity of $A$. This operation is also called *Sampled Dense Dense Matrix Multiplication (SDDMM)*.\n",
|
|
"\n",
|
|
"Enjoying the [batched SDDMM API](https://docs.dgl.ai/en/latest/generated/dgl.sparse.bsddmm.html) in DGL, we can parallel the computation on multiple attention heads (different representation subspaces).\n",
|
|
"\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "dh7zc5v0-JKb"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import dgl\n",
|
|
"import dgl.nn as dglnn\n",
|
|
"import dgl.sparse as dglsp\n",
|
|
"import torch\n",
|
|
"import torch.nn as nn\n",
|
|
"import torch.nn.functional as F\n",
|
|
"import torch.optim as optim\n",
|
|
"\n",
|
|
"from dgl.data import AsGraphPredDataset\n",
|
|
"from dgl.dataloading import GraphDataLoader\n",
|
|
"from ogb.graphproppred import collate_dgl, DglGraphPropPredDataset, Evaluator\n",
|
|
"from ogb.graphproppred.mol_encoder import AtomEncoder\n",
|
|
"from tqdm import tqdm\n",
|
|
"\n",
|
|
"\n",
|
|
"class SparseMHA(nn.Module):\n",
|
|
" \"\"\"Sparse Multi-head Attention Module\"\"\"\n",
|
|
"\n",
|
|
" def __init__(self, hidden_size=80, num_heads=8):\n",
|
|
" super().__init__()\n",
|
|
" self.hidden_size = hidden_size\n",
|
|
" self.num_heads = num_heads\n",
|
|
" self.head_dim = hidden_size // num_heads\n",
|
|
" self.scaling = self.head_dim**-0.5\n",
|
|
"\n",
|
|
" self.q_proj = nn.Linear(hidden_size, hidden_size)\n",
|
|
" self.k_proj = nn.Linear(hidden_size, hidden_size)\n",
|
|
" self.v_proj = nn.Linear(hidden_size, hidden_size)\n",
|
|
" self.out_proj = nn.Linear(hidden_size, hidden_size)\n",
|
|
"\n",
|
|
" def forward(self, A, h):\n",
|
|
" N = len(h)\n",
|
|
" # [N, dh, nh]\n",
|
|
" q = self.q_proj(h).reshape(N, self.head_dim, self.num_heads)\n",
|
|
" q *= self.scaling\n",
|
|
" # [N, dh, nh]\n",
|
|
" k = self.k_proj(h).reshape(N, self.head_dim, self.num_heads)\n",
|
|
" # [N, dh, nh]\n",
|
|
" v = self.v_proj(h).reshape(N, self.head_dim, self.num_heads)\n",
|
|
"\n",
|
|
" ######################################################################\n",
|
|
" # (HIGHLIGHT) Compute the multi-head attention with Sparse Matrix API\n",
|
|
" ######################################################################\n",
|
|
" attn = dglsp.bsddmm(A, q, k.transpose(1, 0)) # (sparse) [N, N, nh]\n",
|
|
" # Sparse softmax by default applies on the last sparse dimension.\n",
|
|
" attn = attn.softmax() # (sparse) [N, N, nh]\n",
|
|
" out = dglsp.bspmm(attn, v) # [N, dh, nh]\n",
|
|
"\n",
|
|
" return self.out_proj(out.reshape(N, -1))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "3_Fm6Lrx-JKc"
|
|
},
|
|
"source": [
|
|
"## Graph Transformer Layer\n",
|
|
"\n",
|
|
"The GT layer is composed of Multi-head Attention, Batch Norm, and Feed-forward Network, connected by residual links as in vanilla transformer.\n",
|
|
"\n",
|
|
"<img src=\"https://drive.google.com/uc?id=1cm-Ijw7bUQIOkoTKn5MQ3m4-66JqCsMz\" width=\"300\">"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "M6h7JVWT-JKd"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"class GTLayer(nn.Module):\n",
|
|
" \"\"\"Graph Transformer Layer\"\"\"\n",
|
|
"\n",
|
|
" def __init__(self, hidden_size=80, num_heads=8):\n",
|
|
" super().__init__()\n",
|
|
" self.MHA = SparseMHA(hidden_size=hidden_size, num_heads=num_heads)\n",
|
|
" self.batchnorm1 = nn.BatchNorm1d(hidden_size)\n",
|
|
" self.batchnorm2 = nn.BatchNorm1d(hidden_size)\n",
|
|
" self.FFN1 = nn.Linear(hidden_size, hidden_size * 2)\n",
|
|
" self.FFN2 = nn.Linear(hidden_size * 2, hidden_size)\n",
|
|
"\n",
|
|
" def forward(self, A, h):\n",
|
|
" h1 = h\n",
|
|
" h = self.MHA(A, h)\n",
|
|
" h = self.batchnorm1(h + h1)\n",
|
|
"\n",
|
|
" h2 = h\n",
|
|
" h = self.FFN2(F.relu(self.FFN1(h)))\n",
|
|
" h = h2 + h\n",
|
|
"\n",
|
|
" return self.batchnorm2(h)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "t40DhVjI-JKd"
|
|
},
|
|
"source": [
|
|
"## Graph Transformer Model\n",
|
|
"\n",
|
|
"The GT model is constructed by stacking GT layers. The input positional encoding of vanilla transformer is replaced with Laplacian positional encoding [(Dwivedi et al. 2020)](https://arxiv.org/abs/2003.00982). For the graph-level prediction task, an extra pooler is stacked on top of GT layers to aggregate node feature of the same graph."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "UrjvEBrF-JKe"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"class GTModel(nn.Module):\n",
|
|
" def __init__(\n",
|
|
" self,\n",
|
|
" out_size,\n",
|
|
" hidden_size=80,\n",
|
|
" pos_enc_size=2,\n",
|
|
" num_layers=8,\n",
|
|
" num_heads=8,\n",
|
|
" ):\n",
|
|
" super().__init__()\n",
|
|
" self.atom_encoder = AtomEncoder(hidden_size)\n",
|
|
" self.pos_linear = nn.Linear(pos_enc_size, hidden_size)\n",
|
|
" self.layers = nn.ModuleList(\n",
|
|
" [GTLayer(hidden_size, num_heads) for _ in range(num_layers)]\n",
|
|
" )\n",
|
|
" self.pooler = dglnn.SumPooling()\n",
|
|
" self.predictor = nn.Sequential(\n",
|
|
" nn.Linear(hidden_size, hidden_size // 2),\n",
|
|
" nn.ReLU(),\n",
|
|
" nn.Linear(hidden_size // 2, hidden_size // 4),\n",
|
|
" nn.ReLU(),\n",
|
|
" nn.Linear(hidden_size // 4, out_size),\n",
|
|
" )\n",
|
|
"\n",
|
|
" def forward(self, g, X, pos_enc):\n",
|
|
" indices = torch.stack(g.edges())\n",
|
|
" N = g.num_nodes()\n",
|
|
" A = dglsp.spmatrix(indices, shape=(N, N))\n",
|
|
" h = self.atom_encoder(X) + self.pos_linear(pos_enc)\n",
|
|
" for layer in self.layers:\n",
|
|
" h = layer(A, h)\n",
|
|
" h = self.pooler(g, h)\n",
|
|
"\n",
|
|
" return self.predictor(h)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"id": "RdrPU18I-JKe"
|
|
},
|
|
"source": [
|
|
"## Training\n",
|
|
"\n",
|
|
"We train the GT model on [ogbg-molhiv](https://ogb.stanford.edu/docs/graphprop/#ogbg-mol) benchmark. The Laplacian positional encoding of each graph is pre-computed (with the API [here](https://docs.dgl.ai/en/latest/generated/dgl.laplacian_pe.html)) as part of the input to the model.\n",
|
|
"\n",
|
|
"*Note that we down-sample the dataset to make this demo runs faster. See the* [*example script*](https://github.com/dmlc/dgl/blob/master/examples/sparse/graph_transformer.py) *for the performance on the full dataset.*"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"id": "V41i0w-9-JKe",
|
|
"colab": {
|
|
"base_uri": "https://localhost:8080/"
|
|
},
|
|
"outputId": "15343d1a-a32d-4677-d053-d9da96910f43"
|
|
},
|
|
"outputs": [
|
|
{
|
|
"output_type": "stream",
|
|
"name": "stderr",
|
|
"text": [
|
|
"Computing Laplacian PE: 1%| | 25/4000 [00:00<00:16, 244.77it/s]/usr/local/lib/python3.8/dist-packages/dgl/backend/pytorch/tensor.py:52: UserWarning: Casting complex values to real discards the imaginary part (Triggered internally at ../aten/src/ATen/native/Copy.cpp:250.)\n",
|
|
" return th.as_tensor(data, dtype=dtype)\n",
|
|
"Computing Laplacian PE: 100%|██████████| 4000/4000 [00:13<00:00, 296.04it/s]\n"
|
|
]
|
|
},
|
|
{
|
|
"output_type": "stream",
|
|
"name": "stdout",
|
|
"text": [
|
|
"Epoch: 000, Loss: 0.2486, Val: 0.3082, Test: 0.3068\n",
|
|
"Epoch: 001, Loss: 0.1695, Val: 0.4684, Test: 0.4572\n",
|
|
"Epoch: 002, Loss: 0.1428, Val: 0.5887, Test: 0.4721\n",
|
|
"Epoch: 003, Loss: 0.1237, Val: 0.6375, Test: 0.5010\n",
|
|
"Epoch: 004, Loss: 0.1127, Val: 0.6628, Test: 0.4854\n",
|
|
"Epoch: 005, Loss: 0.1047, Val: 0.6811, Test: 0.4983\n",
|
|
"Epoch: 006, Loss: 0.0949, Val: 0.6751, Test: 0.5409\n",
|
|
"Epoch: 007, Loss: 0.0901, Val: 0.6340, Test: 0.5357\n",
|
|
"Epoch: 008, Loss: 0.0811, Val: 0.6717, Test: 0.5543\n",
|
|
"Epoch: 009, Loss: 0.0643, Val: 0.7861, Test: 0.5628\n",
|
|
"Epoch: 010, Loss: 0.0489, Val: 0.7319, Test: 0.5341\n",
|
|
"Epoch: 011, Loss: 0.0340, Val: 0.7884, Test: 0.5299\n",
|
|
"Epoch: 012, Loss: 0.0285, Val: 0.5887, Test: 0.4293\n",
|
|
"Epoch: 013, Loss: 0.0361, Val: 0.5514, Test: 0.3419\n",
|
|
"Epoch: 014, Loss: 0.0451, Val: 0.6795, Test: 0.4964\n",
|
|
"Epoch: 015, Loss: 0.0429, Val: 0.7405, Test: 0.5527\n",
|
|
"Epoch: 016, Loss: 0.0331, Val: 0.7859, Test: 0.4994\n",
|
|
"Epoch: 017, Loss: 0.0177, Val: 0.6544, Test: 0.4457\n",
|
|
"Epoch: 018, Loss: 0.0201, Val: 0.8250, Test: 0.6073\n",
|
|
"Epoch: 019, Loss: 0.0093, Val: 0.7356, Test: 0.5561\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"@torch.no_grad()\n",
|
|
"def evaluate(model, dataloader, evaluator, device):\n",
|
|
" model.eval()\n",
|
|
" y_true = []\n",
|
|
" y_pred = []\n",
|
|
" for batched_g, labels in dataloader:\n",
|
|
" batched_g, labels = batched_g.to(device), labels.to(device)\n",
|
|
" y_hat = model(batched_g, batched_g.ndata[\"feat\"], batched_g.ndata[\"PE\"])\n",
|
|
" y_true.append(labels.view(y_hat.shape).detach().cpu())\n",
|
|
" y_pred.append(y_hat.detach().cpu())\n",
|
|
" y_true = torch.cat(y_true, dim=0).numpy()\n",
|
|
" y_pred = torch.cat(y_pred, dim=0).numpy()\n",
|
|
" input_dict = {\"y_true\": y_true, \"y_pred\": y_pred}\n",
|
|
" return evaluator.eval(input_dict)[\"rocauc\"]\n",
|
|
"\n",
|
|
"\n",
|
|
"def train(model, dataset, evaluator, device):\n",
|
|
" train_dataloader = GraphDataLoader(\n",
|
|
" dataset[dataset.train_idx],\n",
|
|
" batch_size=256,\n",
|
|
" shuffle=True,\n",
|
|
" collate_fn=collate_dgl,\n",
|
|
" )\n",
|
|
" valid_dataloader = GraphDataLoader(\n",
|
|
" dataset[dataset.val_idx], batch_size=256, collate_fn=collate_dgl\n",
|
|
" )\n",
|
|
" test_dataloader = GraphDataLoader(\n",
|
|
" dataset[dataset.test_idx], batch_size=256, collate_fn=collate_dgl\n",
|
|
" )\n",
|
|
" optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
|
|
" num_epochs = 20\n",
|
|
" scheduler = optim.lr_scheduler.StepLR(\n",
|
|
" optimizer, step_size=num_epochs, gamma=0.5\n",
|
|
" )\n",
|
|
" loss_fcn = nn.BCEWithLogitsLoss()\n",
|
|
"\n",
|
|
" for epoch in range(num_epochs):\n",
|
|
" model.train()\n",
|
|
" total_loss = 0.0\n",
|
|
" for batched_g, labels in train_dataloader:\n",
|
|
" batched_g, labels = batched_g.to(device), labels.to(device)\n",
|
|
" logits = model(\n",
|
|
" batched_g, batched_g.ndata[\"feat\"], batched_g.ndata[\"PE\"]\n",
|
|
" )\n",
|
|
" loss = loss_fcn(logits, labels.float())\n",
|
|
" total_loss += loss.item()\n",
|
|
" optimizer.zero_grad()\n",
|
|
" loss.backward()\n",
|
|
" optimizer.step()\n",
|
|
" scheduler.step()\n",
|
|
" avg_loss = total_loss / len(train_dataloader)\n",
|
|
" val_metric = evaluate(model, valid_dataloader, evaluator, device)\n",
|
|
" test_metric = evaluate(model, test_dataloader, evaluator, device)\n",
|
|
" print(\n",
|
|
" f\"Epoch: {epoch:03d}, Loss: {avg_loss:.4f}, \"\n",
|
|
" f\"Val: {val_metric:.4f}, Test: {test_metric:.4f}\"\n",
|
|
" )\n",
|
|
"\n",
|
|
"\n",
|
|
"# Training device.\n",
|
|
"dev = torch.device(\"cpu\")\n",
|
|
"# Uncomment the code below to train on GPU. Be sure to install DGL with CUDA support.\n",
|
|
"#dev = torch.device(\"cuda:0\")\n",
|
|
"\n",
|
|
"# Load dataset.\n",
|
|
"pos_enc_size = 8\n",
|
|
"dataset = AsGraphPredDataset(\n",
|
|
" DglGraphPropPredDataset(\"ogbg-molhiv\", \"./data/OGB\")\n",
|
|
")\n",
|
|
"evaluator = Evaluator(\"ogbg-molhiv\")\n",
|
|
"\n",
|
|
"# Down sample the dataset to make the tutorial run faster.\n",
|
|
"import random\n",
|
|
"random.seed(42)\n",
|
|
"train_size = len(dataset.train_idx)\n",
|
|
"val_size = len(dataset.val_idx)\n",
|
|
"test_size = len(dataset.test_idx)\n",
|
|
"dataset.train_idx = dataset.train_idx[\n",
|
|
" torch.LongTensor(random.sample(range(train_size), 2000))\n",
|
|
"]\n",
|
|
"dataset.val_idx = dataset.val_idx[\n",
|
|
" torch.LongTensor(random.sample(range(val_size), 1000))\n",
|
|
"]\n",
|
|
"dataset.test_idx = dataset.test_idx[\n",
|
|
" torch.LongTensor(random.sample(range(test_size), 1000))\n",
|
|
"]\n",
|
|
"\n",
|
|
"# Laplacian positional encoding.\n",
|
|
"indices = torch.cat([dataset.train_idx, dataset.val_idx, dataset.test_idx])\n",
|
|
"for idx in tqdm(indices, desc=\"Computing Laplacian PE\"):\n",
|
|
" g, _ = dataset[idx]\n",
|
|
" g.ndata[\"PE\"] = dgl.laplacian_pe(g, k=pos_enc_size, padding=True)\n",
|
|
"\n",
|
|
"# Create model.\n",
|
|
"out_size = dataset.num_tasks\n",
|
|
"model = GTModel(out_size=out_size, pos_enc_size=pos_enc_size).to(dev)\n",
|
|
"\n",
|
|
"# Kick off training.\n",
|
|
"train(model, dataset, evaluator, dev)"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"language_info": {
|
|
"name": "python"
|
|
},
|
|
"orig_nbformat": 4,
|
|
"colab": {
|
|
"provenance": []
|
|
},
|
|
"gpuClass": "standard",
|
|
"kernelspec": {
|
|
"name": "python3",
|
|
"display_name": "Python 3"
|
|
},
|
|
"accelerator": "GPU"
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|