{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# 构建图卷积网络\n", "\n", "**原作者**: [Yulun Yao](https://yulunyao.io/)，[Chien-Yu Lin](https://homes.cs.washington.edu/~cyulin/)\n", "\n", "本文是介绍性教程，介绍如何使用 Relay 构建图卷积网络（Graph Convolutional Network，简称 GCN）。在本教程中，将在 Cora 数据集上演示 GCN。Cora 数据集是图神经网络（Graph Neural Networks，简称 GNN）和支持 GNN 训练和推理的框架的通用基准。直接从 DGL 库加载数据集，以便与 DGL 进行苹果对苹果的比较。\n", "\n", "DGL 安装请参阅 \n", "\n", "- [DGL 文件](https://docs.dgl.ai/install/index.html)\n", "- [PyTorch 安装指南](https://pytorch.org/get-started/locally/)\n", "\n", "## 用 PyTorch 后端在 DGL 中定义 GCN\n", "\n", "[DGL 示例](https://github.com/dmlc/dgl/tree/master/examples/pytorch/gcn) 部分重用了上面示例中的代码。" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import torch\n", "from torch import nn\n", "from torch.nn import functional as F\n", "import dgl\n", "import networkx as nx\n", "from dgl.nn.pytorch import GraphConv\n", "\n", "\n", "class GCN(nn.Module):\n", " def __init__(self, g, n_infeat, n_hidden, n_classes, n_layers, activation):\n", " super().__init__()\n", " self.g = g\n", " self.layers = nn.ModuleList()\n", " self.layers.append(GraphConv(n_infeat, n_hidden, activation=activation))\n", " for i in range(n_layers - 1):\n", " self.layers.append(GraphConv(n_hidden, n_hidden, activation=activation))\n", " self.layers.append(GraphConv(n_hidden, n_classes))\n", "\n", " def forward(self, features):\n", " h = features\n", " for i, layer in enumerate(self.layers):\n", " # 处理不同 DGL 版本的 api 变更\n", " if dgl.__version__ > \"0.3\":\n", " h = layer(self.g, h)\n", " else:\n", " h = layer(h, self.g)\n", " return h" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 定义加载数据集和评估准确性的函数\n", "\n", "你可以用你自己的数据集代替这一部分，这里我们从 DGL 加载数据：" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from dgl.data import load_data\n", "from collections import namedtuple\n", "\n", "\n", "def load_dataset(dataset=\"cora\"):\n", " args = namedtuple(\"args\", [\"dataset\"])\n", " data = load_data(args(dataset))\n", "\n", " # 删除自循环，以避免重复传递节点的特性给自身\n", " g = data.graph\n", " g.remove_edges_from(nx.selfloop_edges(g))\n", " g.add_edges_from(zip(g.nodes, g.nodes))\n", "\n", " return g, data\n", "\n", "\n", "def evaluate(data, logits):\n", " # 训练阶段中不包含的测试集\n", " test_mask = data.test_mask\n", "\n", " pred = logits.argmax(axis=1)\n", " acc = ((pred == data.labels) * test_mask).sum() / test_mask.sum()\n", "\n", " return acc" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 加载数据并设置模型参数" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " NumNodes: 2708\n", " NumEdges: 10556\n", " NumFeats: 1433\n", " NumClasses: 7\n", " NumTrainingSamples: 140\n", " NumValidationSamples: 500\n", " NumTestSamples: 1000\n", "Done loading data from cached files.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/dgl/data/utils.py:288: UserWarning: Property dataset.graph will be deprecated, please use dataset[0] instead.\n", " warnings.warn('Property {} will be deprecated, please use {} instead.'.format(old, new))\n", "/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/dgl/data/utils.py:288: UserWarning: Property dataset.feat will be deprecated, please use g.ndata['feat'] instead.\n", " warnings.warn('Property {} will be deprecated, please use {} instead.'.format(old, new))\n", "/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/dgl/data/utils.py:288: UserWarning: Property dataset.num_labels will be deprecated, please use dataset.num_classes instead.\n", " warnings.warn('Property {} will be deprecated, please use {} instead.'.format(old, new))\n" ] } ], "source": [ "\"\"\"\n", "Parameters\n", "----------\n", "dataset: str\n", " Name of dataset. You can choose from ['cora', 'citeseer', 'pubmed'].\n", "\n", "num_layer: int\n", " number of hidden layers\n", "\n", "num_hidden: int\n", " number of the hidden units in the hidden layer\n", "\n", "infeat_dim: int\n", " dimension of the input features\n", "\n", "num_classes: int\n", " dimension of model output (Number of classes)\n", "\"\"\"\n", "dataset = \"cora\"\n", "g, data = load_dataset(dataset)\n", "\n", "num_layers = 1\n", "num_hidden = 16\n", "infeat_dim = data.features.shape[1]\n", "num_classes = data.num_labels" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 设定 DGL-PyTorch 模型并获得黄金结果\n", "\n", "被训练的 [weights](https://github.com/dmlc/dgl/blob/master/examples/pytorch/gcn/train.py)。" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "/media/pc/data/4tb/lxw/books/tvm/xinetzone/src\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/dgl/data/utils.py:288: UserWarning: Property dataset.feat will be deprecated, please use g.ndata['feat'] instead.\n", " warnings.warn('Property {} will be deprecated, please use {} instead.'.format(old, new))\n", "/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/dgl/heterograph.py:72: DGLWarning: Recommend creating graphs by `dgl.graph(data)` instead of `dgl.DGLGraph(data)`.\n", " dgl_warning('Recommend creating graphs by `dgl.graph(data)`'\n" ] }, { "data": { "text/plain": [ "" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import env\n", "\n", "from tvm.contrib.download import download_testdata\n", "from dgl import DGLGraph\n", "\n", "features = torch.FloatTensor(data.features)\n", "dgl_g = DGLGraph(g)\n", "\n", "torch_model = GCN(dgl_g, infeat_dim, num_hidden, num_classes, num_layers, F.relu)\n", "\n", "# Download the pretrained weights\n", "model_url = \"https://homes.cs.washington.edu/~cyulin/media/gnn_model/gcn_%s.torch\" % (dataset)\n", "model_path = download_testdata(model_url, \"gcn_%s.pickle\" % (dataset), module=\"gcn_model\")\n", "\n", "# Load the weights into the model\n", "torch_model.load_state_dict(torch.load(model_path))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 运行 DGL 模型并测试其准确性" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Print the first five outputs from DGL-PyTorch execution\n", " tensor([[ 0.2640, -1.0674, 0.0736, 0.7828, -0.7666, -0.0291, -0.1403],\n", " [ 0.2670, -0.9722, 0.0714, 0.6953, -0.6088, -0.0735, -0.1660],\n", " [ 0.2985, -0.9762, 0.1139, 0.5794, -0.5615, -0.0353, -0.1830],\n", " [ 0.2773, -1.2461, 0.0398, 0.9599, -1.0011, 0.0598, -0.1064],\n", " [ 0.3692, -1.0940, 0.0363, 0.6424, -0.6491, 0.0804, -0.1536]])\n", "Test accuracy of DGL results: 5.30%\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/dgl/data/utils.py:288: UserWarning: Property dataset.test_mask will be deprecated, please use g.ndata['test_mask'] instead.\n", " warnings.warn('Property {} will be deprecated, please use {} instead.'.format(old, new))\n", "/media/workspace/anaconda3/envs/torchx/lib/python3.10/site-packages/dgl/data/utils.py:288: UserWarning: Property dataset.label will be deprecated, please use g.ndata['label'] instead.\n", " warnings.warn('Property {} will be deprecated, please use {} instead.'.format(old, new))\n" ] } ], "source": [ "torch_model.eval()\n", "with torch.no_grad():\n", " logits_torch = torch_model(features)\n", "print(\"Print the first five outputs from DGL-PyTorch execution\\n\", logits_torch[:5])\n", "\n", "acc = evaluate(data, logits_torch.numpy())\n", "print(\"Test accuracy of DGL results: {:.2%}\".format(acc))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 在 Relay 中定义图卷积层\n", "\n", "要在 TVM 上运行 GCN，首先需要实现 Graph Convolution Layer。可以参考 [在 DGL 中使用 MXNet 后端实现的 GraphConv 层](https://github.com/dmlc/dgl/blob/master/python/dgl/nn/mxnet/conv/graphconv.py)。\n", "\n", "该层定义如下运算，注意应用了两次转置来保持 sparse_dense 算子右手边的邻接矩阵，这个方法是临时的，在接下来的几周当有稀疏矩阵转置并且支持左稀疏算子的时候会更新。\n", "\n", "```{math}\n", "\\mbox{GraphConv}(A, H, W) = A * H * W\n", " = ((H * W)^t * A^t)^t\n", " = ((W^t * H^t) * A^t)^t\n", "```\n" ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from tvm import relay\n", "from tvm.contrib import graph_executor\n", "import tvm\n", "from tvm import te\n", "\n", "\n", "def GraphConv(layer_name, input_dim, output_dim, adj, input, norm=None, bias=True, activation=None):\n", " \"\"\"\n", " Parameters\n", " ----------\n", " layer_name: str\n", " Name of layer\n", "\n", " input_dim: int\n", " Input dimension per node feature\n", "\n", " output_dim: int,\n", " Output dimension per node feature\n", "\n", " adj: namedtuple,\n", " Graph representation (Adjacency Matrix) in Sparse Format (`data`, `indices`, `indptr`),\n", " where `data` has shape [num_nonzeros], indices` has shape [num_nonzeros], `indptr` has shape [num_nodes + 1]\n", "\n", " input: relay.Expr,\n", " Input feature to current layer with shape [num_nodes, input_dim]\n", "\n", " norm: relay.Expr,\n", " Norm passed to this layer to normalize features before and after Convolution.\n", "\n", " bias: bool\n", " Set bias to True to add bias when doing GCN layer\n", "\n", " activation: ,\n", " Activation function applies to the output. e.g. relay.nn.{relu, sigmoid, log_softmax, softmax, leaky_relu}\n", "\n", " Returns\n", " ----------\n", " output: tvm.relay.Expr\n", " The Output Tensor for this layer [num_nodes, output_dim]\n", " \"\"\"\n", " if norm is not None:\n", " input = relay.multiply(input, norm)\n", "\n", " weight = relay.var(layer_name + \".weight\", shape=(input_dim, output_dim))\n", " weight_t = relay.transpose(weight)\n", " dense = relay.nn.dense(weight_t, input)\n", " output = relay.nn.sparse_dense(dense, adj)\n", " output_t = relay.transpose(output)\n", " if norm is not None:\n", " output_t = relay.multiply(output_t, norm)\n", " if bias is True:\n", " _bias = relay.var(layer_name + \".bias\", shape=(output_dim, 1))\n", " output_t = relay.nn.bias_add(output_t, _bias, axis=-1)\n", " if activation is not None:\n", " output_t = activation(output_t)\n", " return output_t" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 准备 GraphConv 层中所需的参数" ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "collapsed": false }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/tmp/ipykernel_523780/3618165654.py:12: DeprecationWarning: \n", "\n", "The scipy.sparse array containers will be used instead of matrices\n", "in Networkx 3.0. Use `to_scipy_sparse_array` instead.\n", " adjacency = nx.to_scipy_sparse_matrix(g)\n" ] } ], "source": [ "import numpy as np\n", "import networkx as nx\n", "\n", "\n", "def prepare_params(g, data):\n", " params = {}\n", " params[\"infeats\"] = data.features.numpy().astype(\n", " \"float32\"\n", " ) # Only support float32 as feature for now\n", "\n", " # Generate adjacency matrix\n", " adjacency = nx.to_scipy_sparse_matrix(g)\n", " params[\"g_data\"] = adjacency.data.astype(\"float32\")\n", " params[\"indices\"] = adjacency.indices.astype(\"int32\")\n", " params[\"indptr\"] = adjacency.indptr.astype(\"int32\")\n", "\n", " # Normalization w.r.t. node degrees\n", " degs = [g.in_degree[i] for i in range(g.number_of_nodes())]\n", " params[\"norm\"] = np.power(degs, -0.5).astype(\"float32\")\n", " params[\"norm\"] = params[\"norm\"].reshape((params[\"norm\"].shape[0], 1))\n", "\n", " return params\n", "\n", "\n", "params = prepare_params(g, data)\n", "\n", "# Check shape of features and the validity of adjacency matrix\n", "assert len(params[\"infeats\"].shape) == 2\n", "assert (\n", " params[\"g_data\"] is not None and params[\"indices\"] is not None and params[\"indptr\"] is not None\n", ")\n", "assert params[\"infeats\"].shape[0] == params[\"indptr\"].shape[0] - 1" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 把层放在一起" ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Define input features, norms, adjacency matrix in Relay\n", "infeats = relay.var(\"infeats\", shape=data.features.shape)\n", "norm = relay.Constant(tvm.nd.array(params[\"norm\"]))\n", "g_data = relay.Constant(tvm.nd.array(params[\"g_data\"]))\n", "indices = relay.Constant(tvm.nd.array(params[\"indices\"]))\n", "indptr = relay.Constant(tvm.nd.array(params[\"indptr\"]))\n", "\n", "Adjacency = namedtuple(\"Adjacency\", [\"data\", \"indices\", \"indptr\"])\n", "adj = Adjacency(g_data, indices, indptr)\n", "\n", "# Construct the 2-layer GCN\n", "layers = []\n", "layers.append(\n", " GraphConv(\n", " layer_name=\"layers.0\",\n", " input_dim=infeat_dim,\n", " output_dim=num_hidden,\n", " adj=adj,\n", " input=infeats,\n", " norm=norm,\n", " activation=relay.nn.relu,\n", " )\n", ")\n", "layers.append(\n", " GraphConv(\n", " layer_name=\"layers.1\",\n", " input_dim=num_hidden,\n", " output_dim=num_classes,\n", " adj=adj,\n", " input=layers[-1],\n", " norm=norm,\n", " activation=None,\n", " )\n", ")\n", "\n", "# Analyze free variables and generate Relay function\n", "output = layers[-1]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 使用 TVM 编译并运行\n", "\n", "从 PyTorch 模型导出权重到 Python Dict：" ] }, { "cell_type": "code", "execution_count": 10, "metadata": { "collapsed": false }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "One or more operators have not been tuned. Please tune your model for better performance. Use DEBUG logging level to see more details.\n" ] } ], "source": [ "model_params = {}\n", "for param_tensor in torch_model.state_dict():\n", " model_params[param_tensor] = torch_model.state_dict()[param_tensor].numpy()\n", "\n", "for i in range(num_layers + 1):\n", " params[\"layers.%d.weight\" % (i)] = model_params[\"layers.%d.weight\" % (i)]\n", " params[\"layers.%d.bias\" % (i)] = model_params[\"layers.%d.bias\" % (i)]\n", "\n", "# Set the TVM build target\n", "target = \"llvm\" # Currently only support `llvm` as target\n", "\n", "func = relay.Function(relay.analysis.free_vars(output), output)\n", "func = relay.build_module.bind_params_by_name(func, params)\n", "mod = tvm.IRModule()\n", "mod[\"main\"] = func\n", "# Build with Relay\n", "with tvm.transform.PassContext(opt_level=0): # Currently only support opt_level=0\n", " lib = relay.build(mod, target, params=params)\n", "\n", "# Generate graph executor\n", "dev = tvm.device(target, 0)\n", "m = graph_executor.GraphModule(lib[\"default\"](dev))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## 运行 TVM 模型，测试准确性并通过 DGL 验证" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "collapsed": false }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Print the first five outputs from TVM execution\n", " [[ 0.26396316 -1.067397 0.07361096 0.78283393 -0.7665647 -0.02912378\n", " -0.14030665]\n", " [ 0.2670483 -0.97222644 0.07140031 0.6953188 -0.60881317 -0.07351625\n", " -0.16601387]\n", " [ 0.29854178 -0.97619903 0.11394241 0.57936156 -0.5615169 -0.03528827\n", " -0.18298927]\n", " [ 0.2773209 -1.2461467 0.0398193 0.95992005 -1.0011221 0.059847\n", " -0.10642916]\n", " [ 0.3691777 -1.0940018 0.03631139 0.6423676 -0.6491406 0.08039594\n", " -0.1535899 ]]\n", "Test accuracy of TVM results: 5.30%\n" ] } ], "source": [ "m.run()\n", "logits_tvm = m.get_output(0).numpy()\n", "print(\"Print the first five outputs from TVM execution\\n\", logits_tvm[:5])\n", "\n", "labels = data.labels\n", "test_mask = data.test_mask\n", "\n", "acc = evaluate(data, logits_tvm)\n", "print(\"Test accuracy of TVM results: {:.2%}\".format(acc))\n", "\n", "import tvm.testing\n", "\n", "# Verify the results with the DGL model\n", "tvm.testing.assert_allclose(logits_torch, logits_tvm, atol=1e-3)" ] } ], "metadata": { "interpreter": { "hash": "7a45eadec1f9f49b0fdfd1bc7d360ac982412448ce738fa321afc640e3212175" }, "kernelspec": { "display_name": "Python 3.10.4 ('torchx')", "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.10.4" } }, "nbformat": 4, "nbformat_minor": 0 }