Colab5.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XuXWJLEm2UWS"
      },
      "source": [
        "# CS224W - Colab 5\n",
        "\n",
        "本实验中我们将使用 `PyG`、`DeepSnap` 和 `NetworkX` 在加大规模的 GNN 上做实验。\n",
        "\n",
        "首先我们使用 `PyG` 的 `NeighborSampler` 来提升在 `arxiv` 数据集上训练和测试的规模。\n",
        "\n",
        "然后使用 `DeepSNAP` 和 `NetworkX` 实现了一个简化版本的 `NeighborSampler` 并在 `Cora` 图上以不同采样率进行了多次实验。\n",
        "\n",
        "最终我们会使用不同的分割算法把 `Cora` 图分割成几个社群，然后使用原始 `Cluster-GCN` 的方式来训练模型。\n",
        "\n",
        "**注意**：确保顺序执行所有代码单元，否则变量或库无法在之后的单元中使用。实验愉快 ：）\n",
        "\n",
        "> 译者注：在原课程中，本次实验被取消了。因此本实验笔记是一个教程，不需要实现任何东西，以供参阅。后文提到实现函数其实都是实现好的。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PRfgbfTjCRD_"
      },
      "outputs": [],
      "source": [
        "import torch_geometric\n",
        "torch_geometric.__version__"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SxkYLgxAOxz7"
      },
      "source": [
        "## 1 `PyG` 的邻居采样\n",
        "\n",
        "在邻居采样在 **GraphSAGE** ([Hamilton et al. (2017)](https://arxiv.org/abs/1706.02216)) 刚提出来时，它就是提升 GNN 规模的代表性的方法。正如 Lecture 中讲的，每次训练的时候，只有**K-跳**领域内的结点会被送到显卡中。为了进一步减少负担，我们可以在领域的结点中采样一个子集，以供 GNN 聚合。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kho6SHUVO1ny"
      },
      "source": [
        "### 配置"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X1WJLGKsOx_k"
      },
      "outputs": [],
      "source": [
        "import copy\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import torch_geometric.transforms as T\n",
        "\n",
        "from torch_geometric.nn import SAGEConv\n",
        "from torch_geometric.loader import NeighborSampler\n",
        "from ogb.nodeproppred import PygNodePropPredDataset, Evaluator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UKqZWqRbO7km"
      },
      "source": [
        "### 邻居采样器\n",
        "\n",
        "`PyG` 已经实现了邻居采样的方法，具体为 `torch_geometric.loader` 中的 [NeighborSampler](https://pytorch-geometric.readthedocs.io/en/latest/modules/loader.html#torch_geometric.loader.NeighborSampler)。以下是一个在 OGB 数据集 `arxiv` 上使用邻居采样的示例。\n",
        "\n",
        "如果你内存高效的聚合感兴趣，可以看看 `PyG` 的 [内存高效聚合(Memory-Efficient Aggregations)](https://pytorch-geometric.readthedocs.io/en/latest/notes/sparse_tensor.html)。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nWlyStlRO6_u"
      },
      "outputs": [],
      "source": [
        "dataset_name = 'ogbn-arxiv'\n",
        "dataset = PygNodePropPredDataset(name=dataset_name,\n",
        "                                 transform=T.ToSparseTensor())\n",
        "data = dataset[0]\n",
        "data.adj_t = data.adj_t.to_symmetric()\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "\n",
        "print(f'Device: {device}')\n",
        "\n",
        "data = data.to(device)\n",
        "split_idx = dataset.get_idx_split()\n",
        "train_idx = split_idx['train'].to(device)\n",
        "\n",
        "# Construct the training dataloader for training data\n",
        "# Sample 10 neighbors for each node in the first layer and 5 for the second layer\n",
        "train_loader = NeighborSampler(data.adj_t, node_idx=train_idx,\n",
        "                               sizes=[10, 5], batch_size=4096,\n",
        "                               shuffle=True, num_workers=2)\n",
        "\n",
        "# Specify size as -1 to include all neighbors\n",
        "all_loader = NeighborSampler(data.adj_t, node_idx=None, sizes=[-1],\n",
        "                                  batch_size=4096, shuffle=False,\n",
        "                                  num_workers=2)\n",
        "evaluator = Evaluator(name='ogbn-arxiv')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VjdkIcFpRYyl"
      },
      "source": [
        "### GNN 模型\n",
        "\n",
        "在创建了 `NeighborSampler` 之后，我们也可以修改模型以支持小批量训练。\n",
        "\n",
        "模型的 `forward` 的参数有结点特征 `x`，一个列表的三元组 `adjs`。`adjs` 中每个三元组都有如下元素：\n",
        "\n",
        "- `edge_index`: 描述源结点和目标结点之间的连通性的张量，可以组成一个二部图。\n",
        "- `e_id`: 在原图中，边的索引。\n",
        "- `size`: 二部图的形状，格式为（*源结点数量*，*目标结点数量*）。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gRBJS_5qRWbu"
      },
      "outputs": [],
      "source": [
        "class SAGE(torch.nn.Module):\n",
        "    def __init__(self, input_dim, hidden_dim, output_dim, num_layers,\n",
        "                 dropout):\n",
        "        super(SAGE, self).__init__()\n",
        "\n",
        "        self.convs = torch.nn.ModuleList()\n",
        "        self.bns = torch.nn.ModuleList()\n",
        "\n",
        "        self.convs.append(SAGEConv(input_dim, hidden_dim))\n",
        "        self.bns.append(torch.nn.BatchNorm1d(hidden_dim))\n",
        "\n",
        "        for i in range(num_layers - 2):\n",
        "            self.convs.append(\n",
        "                SAGEConv(hidden_dim, hidden_dim))\n",
        "            self.bns.append(torch.nn.BatchNorm1d(hidden_dim))\n",
        "        self.convs.append(SAGEConv(hidden_dim, output_dim))\n",
        "\n",
        "        self.softmax = torch.nn.LogSoftmax(dim=1)\n",
        "\n",
        "        self.dropout = dropout\n",
        "\n",
        "        self.num_layers = num_layers\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        for conv in self.convs:\n",
        "            conv.reset_parameters()\n",
        "        for bn in self.bns:\n",
        "            bn.reset_parameters()\n",
        "\n",
        "    def forward(self, x, adjs, mode=\"batch\"):\n",
        "        if mode == \"batch\":\n",
        "            for i, (edge_index, _, size) in enumerate(adjs):\n",
        "                # Extract target node features\n",
        "                x_target = x[:size[1]]\n",
        "\n",
        "                # Update x for next layer reuse\n",
        "                x = self.convs[i]((x, x_target), edge_index)\n",
        "                if i != self.num_layers - 1:\n",
        "                    x = self.bns[i](x)\n",
        "                    x = F.relu(x)\n",
        "                    x = F.dropout(x, p=self.dropout, training=self.training)\n",
        "        else:\n",
        "            for i, conv in enumerate(self.convs):\n",
        "                x = conv(x, adjs)\n",
        "                if i != self.num_layers - 1:\n",
        "                    x = self.bns[i](x)\n",
        "                    x = F.relu(x)\n",
        "                    x = F.dropout(x, p=self.dropout, training=self.training)\n",
        "        return self.softmax(x)\n",
        "    \n",
        "    def inference(self, x_all, all_loader):\n",
        "        # This function will be called in test\n",
        "        for i in range(self.num_layers):\n",
        "            xs = []\n",
        "            for batch_size, n_id, adj in all_loader:\n",
        "                edge_index, _, size = adj.to(device)\n",
        "                x = x_all[n_id].to(device)\n",
        "                x_target = x[:size[1]]\n",
        "                x = self.convs[i]((x, x_target), edge_index)\n",
        "                if i != self.num_layers - 1:\n",
        "                    x = self.bns[i](x)\n",
        "                    x = F.relu(x)\n",
        "                    x = F.dropout(x, p=self.dropout, training=self.training)\n",
        "                \n",
        "                # Append the node embeddings to xs\n",
        "                xs.append(x.cpu())\n",
        "            \n",
        "            # Concat all embeddings into one tensor\n",
        "            x_all = torch.cat(xs, dim=0)\n",
        "\n",
        "        return x_all"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7cfm7K3wRqqY"
      },
      "source": [
        "### 训练和测试\n",
        "\n",
        "现在我们来实现训练和测试函数。\n",
        "\n",
        "在训练和测试中，我们都需要从 `dataloader` 中采样 `batch`。\n",
        "\n",
        "`NeighborSampler` 的 `dataloader` 中每一个 `batch` 都有三个元素：\n",
        "\n",
        "- `batch_size`: `dataloader` 中指定的 `batch` 的大小。\n",
        "- `n_id`: 邻接矩阵中使用的所有结点（索引的形式）。\n",
        "- `adjs`: 三元组。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-JN0-_QCRn8N"
      },
      "outputs": [],
      "source": [
        "def train(model, data, train_loader, train_idx, optimizer, loss_fn, mode=\"batch\"):\n",
        "    model.train()\n",
        "\n",
        "    total_loss = 0\n",
        "    if mode == \"batch\":\n",
        "        for batch_size, n_id, adjs in train_loader:\n",
        "            # Move all adj sparse tensors to GPU\n",
        "            adjs = [adj.to(device) for adj in adjs]\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            # Index on the node features\n",
        "            out = model(data.x[n_id], adjs)\n",
        "            train_label = data.y[n_id[:batch_size]].squeeze(-1)\n",
        "            loss = loss_fn(out, train_label)\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "            total_loss += loss.item()\n",
        "    else:\n",
        "        optimizer.zero_grad()\n",
        "        out = model(data.x, data.adj_t, mode=mode)[train_idx]\n",
        "        train_label = data.y.squeeze(1)[train_idx]\n",
        "        loss = loss_fn(out, train_label)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        total_loss = loss.item()\n",
        "\n",
        "    return total_loss\n",
        "\n",
        "@torch.no_grad()\n",
        "def test(model, data, all_loader, split_idx, evaluator, mode=\"batch\"):\n",
        "    model.eval()\n",
        "\n",
        "    if mode == \"batch\":\n",
        "        out = model.inference(data.x, all_loader)\n",
        "    else:\n",
        "        out = model(data.x, data.adj_t, mode=\"all\")\n",
        "\n",
        "    y_true = data.y.cpu()\n",
        "    y_pred = out.argmax(dim=-1, keepdim=True)\n",
        "\n",
        "    train_acc = evaluator.eval({\n",
        "        'y_true': y_true[split_idx['train']],\n",
        "        'y_pred': y_pred[split_idx['train']],\n",
        "    })['acc']\n",
        "    valid_acc = evaluator.eval({\n",
        "        'y_true': y_true[split_idx['valid']],\n",
        "        'y_pred': y_pred[split_idx['valid']],\n",
        "    })['acc']\n",
        "    test_acc = evaluator.eval({\n",
        "        'y_true': y_true[split_idx['test']],\n",
        "        'y_pred': y_pred[split_idx['test']],\n",
        "    })['acc']\n",
        "\n",
        "    return train_acc, valid_acc, test_acc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AiehZ8OiR2q9"
      },
      "source": [
        "### 小批量训练"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zFaI2eCARy0v"
      },
      "outputs": [],
      "source": [
        "args = {\n",
        "    'device': device,\n",
        "    'num_layers': 2,\n",
        "    'hidden_dim': 128,\n",
        "    'dropout': 0.5,\n",
        "    'lr': 0.01,\n",
        "    'epochs': 100,\n",
        "}\n",
        "\n",
        "batch_model = SAGE(data.num_features, args['hidden_dim'],\n",
        "            dataset.num_classes, args['num_layers'],\n",
        "            args['dropout']).to(device)\n",
        "batch_model.reset_parameters()\n",
        "\n",
        "optimizer = torch.optim.Adam(batch_model.parameters(), lr=args['lr'])\n",
        "loss_fn = F.nll_loss\n",
        "\n",
        "best_batch_model = None\n",
        "best_valid_acc = 0\n",
        "\n",
        "batch_results = []\n",
        "\n",
        "for epoch in range(1, 1 + args[\"epochs\"]):\n",
        "    loss = train(batch_model, data, train_loader, train_idx, optimizer, loss_fn, mode=\"batch\")\n",
        "    result = test(batch_model, data, all_loader, split_idx, evaluator, mode=\"batch\")\n",
        "    batch_results.append(result)\n",
        "    train_acc, valid_acc, test_acc = result\n",
        "    if valid_acc > best_valid_acc:\n",
        "        best_valid_acc = valid_acc\n",
        "        best_batch_model = copy.deepcopy(batch_model)\n",
        "    print(f'Epoch: {epoch:02d}, '\n",
        "          f'Loss: {loss:.4f}, '\n",
        "          f'Train: {100 * train_acc:.2f}%, '\n",
        "          f'Valid: {100 * valid_acc:.2f}% '\n",
        "          f'Test: {100 * test_acc:.2f}%')\n",
        "best_result = test(best_batch_model, data, all_loader, split_idx, evaluator, mode=\"batch\")\n",
        "train_acc, valid_acc, test_acc = best_result\n",
        "print(f'Best model: '\n",
        "      f'Train: {100 * train_acc:.2f}%, '\n",
        "      f'Valid: {100 * valid_acc:.2f}% '\n",
        "      f'Test: {100 * test_acc:.2f}%')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-OyqW-1pSMLW"
      },
      "source": [
        "### 全批量训练"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mU5eAviTSFMO"
      },
      "outputs": [],
      "source": [
        "# Use the same parameters for a full-batch training\n",
        "args = {\n",
        "    'device': device,\n",
        "    'num_layers': 2,\n",
        "    'hidden_dim': 128,\n",
        "    'dropout': 0.5,\n",
        "    'lr': 0.01,\n",
        "    'epochs': 100,\n",
        "}\n",
        "\n",
        "all_model = SAGE(data.num_features, args['hidden_dim'],\n",
        "            dataset.num_classes, args['num_layers'],\n",
        "            args['dropout']).to(device)\n",
        "all_model.reset_parameters()\n",
        "\n",
        "optimizer = torch.optim.Adam(all_model.parameters(), lr=args['lr'])\n",
        "loss_fn = F.nll_loss\n",
        "\n",
        "best_all_model = None\n",
        "best_valid_acc = 0\n",
        "\n",
        "all_results = []\n",
        "\n",
        "for epoch in range(1, 1 + args[\"epochs\"]):\n",
        "    loss = train(all_model, data, train_loader, train_idx, optimizer, loss_fn, mode=\"all\")\n",
        "    result = test(all_model, data, all_loader, split_idx, evaluator, mode=\"all\")\n",
        "    all_results.append(result)\n",
        "    train_acc, valid_acc, test_acc = result\n",
        "    if valid_acc > best_valid_acc:\n",
        "        best_valid_acc = valid_acc\n",
        "        best_all_model = copy.deepcopy(all_model)\n",
        "    print(f'Epoch: {epoch:02d}, '\n",
        "          f'Loss: {loss:.4f}, '\n",
        "          f'Train: {100 * train_acc:.2f}%, '\n",
        "          f'Valid: {100 * valid_acc:.2f}% '\n",
        "          f'Test: {100 * test_acc:.2f}%')\n",
        "best_result = test(best_all_model, data, all_loader, split_idx, evaluator, mode=\"all\")\n",
        "train_acc, valid_acc, test_acc = best_result\n",
        "print(f'Best model: '\n",
        "      f'Train: {100 * train_acc:.2f}%, '\n",
        "      f'Valid: {100 * valid_acc:.2f}% '\n",
        "      f'Test: {100 * test_acc:.2f}%')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NrECcOQQSZo1"
      },
      "source": [
        "### 可视化"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sh_qvSG1SV63"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "batch_results = np.array(batch_results)\n",
        "all_results = np.array(all_results)\n",
        "\n",
        "x = np.arange(1, 101)\n",
        "\n",
        "plt.figure(figsize=(9, 7))\n",
        "\n",
        "plt.plot(x, batch_results[:, 1], label=\"Batch Validation\")\n",
        "plt.plot(x, batch_results[:, 2], label=\"Batch Test\")\n",
        "plt.plot(x, all_results[:, 1], label=\"All Validation\")\n",
        "plt.plot(x, all_results[:, 2], label=\"All Test\")\n",
        "plt.title('Model Accuracy')\n",
        "plt.ylabel('Accuracy')\n",
        "plt.xlabel('Epoch')\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WFb2OAvOSn_O"
      },
      "source": [
        "## 2 不同比率的邻居采样\n",
        "\n",
        "现在我们使用 `DeepSNAP` 和 `NetworkX` 来实现一个简化版的邻居采样，并用不同比率的邻域采样来训练模型。\n",
        "\n",
        "为了让实验进行的更快，我们这里使用 `Cora` 图。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P9U0F7bnSz9u"
      },
      "source": [
        "### 配置"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PUF4on-fSxcq"
      },
      "outputs": [],
      "source": [
        "import copy\n",
        "import torch\n",
        "import random\n",
        "import numpy as np\n",
        "import networkx as nx\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from torch_geometric.nn import SAGEConv\n",
        "from torch.utils.data import DataLoader\n",
        "from torch_geometric.datasets import Planetoid\n",
        "from torch.nn import Sequential, Linear, ReLU\n",
        "from deepsnap.dataset import GraphDataset\n",
        "from deepsnap.graph import Graph\n",
        "\n",
        "pyg_dataset = Planetoid('./dataset', \"Cora\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qw6k-KdFTEYw"
      },
      "source": [
        "### GNN 模型"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PvUlNi2TS09i"
      },
      "outputs": [],
      "source": [
        "class GNN(torch.nn.Module):\n",
        "    def __init__(self, input_dim, hidden_dim, output_dim, args):\n",
        "        super(GNN, self).__init__()\n",
        "        self.dropout = args['dropout']\n",
        "        self.num_layers = args['num_layers']\n",
        "\n",
        "        self.convs = nn.ModuleList()\n",
        "        self.bns = nn.ModuleList()\n",
        "\n",
        "        self.convs.append(SAGEConv(input_dim, hidden_dim))\n",
        "        self.bns.append(nn.BatchNorm1d(hidden_dim))\n",
        "\n",
        "        for l in range(self.num_layers - 2):\n",
        "            self.convs.append(SAGEConv(hidden_dim, hidden_dim))\n",
        "            self.bns.append(nn.BatchNorm1d(hidden_dim))\n",
        "        self.convs.append(SAGEConv(hidden_dim, hidden_dim))\n",
        "\n",
        "        self.post_mp = nn.Linear(hidden_dim, output_dim)\n",
        "\n",
        "    def forward(self, data, mode=\"batch\"):\n",
        "        if mode == \"batch\":\n",
        "            edge_indices, x = data\n",
        "            for i in range(len(self.convs) - 1):\n",
        "                edge_index = edge_indices[i]\n",
        "                x = self.convs[i](x, edge_index)\n",
        "                x = self.bns[i](x)\n",
        "                x = F.relu(x)\n",
        "                x = F.dropout(x, p=self.dropout, training=self.training)\n",
        "            x = self.convs[-1](x, edge_indices[len(self.convs) - 1])\n",
        "        else:\n",
        "            x, edge_index = data.node_feature, data.edge_index\n",
        "            for i in range(len(self.convs) - 1):\n",
        "                x = self.convs[i](x, edge_index)\n",
        "                x = self.bns[i](x)\n",
        "                x = F.relu(x)\n",
        "                x = F.dropout(x, p=self.dropout, training=self.training)\n",
        "            x = self.convs[-1](x, edge_index)\n",
        "        x = self.post_mp(x)\n",
        "        x = F.log_softmax(x, dim=1)\n",
        "        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ulp1A3evcJ-I"
      },
      "source": [
        "### 邻居采样\n",
        "\n",
        "我们使用 `DeepSNAP` 和 `NetworkX` 实现了能够采样邻居的函数。\n",
        "\n",
        "注意 `Cora` 上的结点分类任务是一个半监督的分类任务，我们这里通过将最后的比率设为 1 来保留所有带标签的结点。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LI4qHkE4cQOh"
      },
      "outputs": [],
      "source": [
        "def sample_neighbors(nodes, G, ratio, all_nodes):\n",
        "    # This fuction takes a set of nodes, a NetworkX graph G and neighbor sampling ratio.\n",
        "    # It will return sampled neighbors (unioned with input nodes) and edges between \n",
        "    neighbors = set()\n",
        "    edges = []\n",
        "    for node in nodes:\n",
        "        neighbors_list = list(nx.neighbors(G, node))\n",
        "\n",
        "        # We only sample the (ratio * number of neighbors) neighbors\n",
        "        num = int(len(neighbors_list) * ratio)\n",
        "        if num > 0:\n",
        "            # Random shuffle the neighbors\n",
        "            random.shuffle(neighbors_list)\n",
        "            neighbors_list = neighbors_list[:num]\n",
        "            for neighbor in neighbors_list:\n",
        "                # Add neighbors\n",
        "                neighbors.add(neighbor)\n",
        "                edges.append((neighbor, node))\n",
        "    return neighbors, neighbors.union(all_nodes), edges\n",
        "\n",
        "def nodes_to_tensor(nodes):\n",
        "    # This function transform a set of nodes to node index tensor\n",
        "    node_label_index = torch.tensor(list(nodes), dtype=torch.long)\n",
        "    return node_label_index\n",
        "\n",
        "def edges_to_tensor(edges):\n",
        "    # This function transform a set of edges to edge index tensor\n",
        "    edge_index = torch.tensor(list(edges), dtype=torch.long)\n",
        "    edge_index = torch.cat([edge_index, torch.flip(edge_index, [1])], dim=0)\n",
        "    edge_index = edge_index.permute(1, 0)\n",
        "    return edge_index\n",
        "\n",
        "def relable(nodes, labeled_nodes, edges_list):\n",
        "    # Relable the nodes, labeled_nodes and edges_list\n",
        "    relabled_edges_list = []\n",
        "    sorted_nodes = sorted(nodes)\n",
        "    node_mapping = {node : i for i, node in enumerate(sorted_nodes)}\n",
        "    for orig_edges in edges_list:\n",
        "        relabeled_edges = []\n",
        "        for edge in orig_edges:\n",
        "            relabeled_edges.append((node_mapping[edge[0]], node_mapping[edge[1]]))\n",
        "        relabled_edges_list.append(relabeled_edges)\n",
        "    relabeled_labeled_nodes = [node_mapping[node] for node in labeled_nodes]\n",
        "    relabeled_nodes = [node_mapping[node] for node in nodes]\n",
        "    return relabled_edges_list, relabeled_nodes, relabeled_labeled_nodes\n",
        "\n",
        "def neighbor_sampling(graph, K=2, ratios=(0.1, 0.1, 0.1)):\n",
        "    # This function takes a DeepSNAP graph, K the number of GNN layers, and neighbor \n",
        "    # sampling ratios for each layer. This function returns relabeled node feature, \n",
        "    # edge indices and node_label_index\n",
        "\n",
        "    assert K + 1 == len(ratios)\n",
        "\n",
        "    labeled_nodes = graph.node_label_index.tolist()\n",
        "    random.shuffle(labeled_nodes)\n",
        "    num = int(len(labeled_nodes) * ratios[-1])\n",
        "    if num > 0:\n",
        "        labeled_nodes = labeled_nodes[:num]\n",
        "    nodes_list = [set(labeled_nodes)]\n",
        "    edges_list = []\n",
        "    all_nodes = labeled_nodes\n",
        "    for k in range(K):\n",
        "        # Get nodes and edges from the previous layer\n",
        "        nodes, all_nodes, edges = \\\n",
        "            sample_neighbors(nodes_list[-1], graph.G, ratios[len(ratios) - k - 2], all_nodes)\n",
        "        nodes_list.append(nodes)\n",
        "        edges_list.append(edges)\n",
        "    \n",
        "    # Reverse the lists\n",
        "    nodes_list.reverse()\n",
        "    edges_list.reverse()\n",
        "\n",
        "    relabled_edges_list, relabeled_all_nodes, relabeled_labeled_nodes = \\\n",
        "        relable(all_nodes, labeled_nodes, edges_list)\n",
        "\n",
        "    node_index = nodes_to_tensor(relabeled_all_nodes)\n",
        "    # All node features that will be used\n",
        "    node_feature = graph.node_feature[node_index]\n",
        "    edge_indices = [edges_to_tensor(edges) for edges in relabled_edges_list]\n",
        "    node_label_index = nodes_to_tensor(relabeled_labeled_nodes)\n",
        "    print(f\"Sampled {node_feature.shape[0]} nodes, {edge_indices[0].shape[1] // 2} edges, {node_label_index.shape[0]} labeled nodes\")\n",
        "    return node_feature, edge_indices, node_label_index"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ooy6Hcf7TIhI"
      },
      "source": [
        "### 训练和测试"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iSmZhpzPTGPY"
      },
      "outputs": [],
      "source": [
        "def train(train_graphs, val_graphs, args, model, optimizer, mode=\"batch\"):\n",
        "    best_val = 0\n",
        "    best_model = None\n",
        "    accs = []\n",
        "    graph_train = train_graphs[0]\n",
        "    graph_train.to(args['device'])\n",
        "    for epoch in range(1, 1 + args['epochs']):\n",
        "        model.train()\n",
        "        optimizer.zero_grad()\n",
        "        if mode == \"batch\":\n",
        "            node_feature, edge_indices, node_label_index = neighbor_sampling(graph_train, args['num_layers'], args['ratios'])\n",
        "            node_feature = node_feature.to(args['device'])\n",
        "            node_label_index = node_label_index.to(args['device'])\n",
        "            for i in range(len(edge_indices)):\n",
        "                edge_indices[i] = edge_indices[i].to(args['device'])\n",
        "            pred = model([edge_indices, node_feature])\n",
        "            pred = pred[node_label_index]\n",
        "            label = graph_train.node_label[node_label_index]\n",
        "        elif mode == \"community\":\n",
        "            graph = random.choice(train_graphs)\n",
        "            graph = graph.to(args['device'])\n",
        "            pred = model(graph, mode=\"all\")\n",
        "            pred = pred[graph.node_label_index]\n",
        "            label = graph.node_label[graph.node_label_index]\n",
        "        else:\n",
        "            pred = model(graph_train, mode=\"all\")\n",
        "            label = graph_train.node_label\n",
        "            pred = pred[graph_train.node_label_index]\n",
        "        loss = F.nll_loss(pred, label)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "        train_acc, val_acc, test_acc = test(val_graphs, model)\n",
        "        accs.append((train_acc, val_acc, test_acc))\n",
        "        if val_acc > best_val:\n",
        "            best_val = val_acc\n",
        "            best_model = copy.deepcopy(model)\n",
        "        print(f'Epoch: {epoch:02d}, '\n",
        "              f'Loss: {loss:.4f}, '\n",
        "              f'Train: {100 * train_acc:.2f}%, '\n",
        "              f'Valid: {100 * val_acc:.2f}% '\n",
        "              f'Test: {100 * test_acc:.2f}%')\n",
        "    return best_model, accs\n",
        "\n",
        "def test(graphs, model):\n",
        "    model.eval()\n",
        "    accs = []\n",
        "    for graph in graphs:\n",
        "        graph = graph.to(args['device'])\n",
        "        pred = model(graph, mode=\"all\")\n",
        "        label = graph.node_label\n",
        "        pred = pred[graph.node_label_index].max(1)[1]\n",
        "        acc = pred.eq(label).sum().item()\n",
        "        acc /= len(label)\n",
        "        accs.append(acc)\n",
        "    return accs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HV7i0v0ETKzf"
      },
      "outputs": [],
      "source": [
        "args = {\n",
        "    'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),\n",
        "    'dropout': 0.5,\n",
        "    'num_layers': 2,\n",
        "    'hidden_size': 64,\n",
        "    'lr': 0.005,\n",
        "    'epochs': 50,\n",
        "    'ratios': (0.8, 0.8, 1),\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rLpRYKbnTQnj"
      },
      "source": [
        "### 全批量训练"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pMGGjbJBTOo1"
      },
      "outputs": [],
      "source": [
        "graphs_train, graphs_val, graphs_test = \\\n",
        "    GraphDataset.pyg_to_graphs(pyg_dataset, verbose=True, fixed_split=True)\n",
        "\n",
        "graph_train = graphs_train[0]\n",
        "graph_val = graphs_val[0]\n",
        "graph_test = graphs_test[0]\n",
        "\n",
        "model = GNN(graph_train.num_node_features, args['hidden_size'], graph_train.num_node_labels, args).to(args['device'])\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'])\n",
        "graphs = [graph_train, graph_val, graph_test]\n",
        "all_best_model, all_accs = train(graphs, graphs, args, model, optimizer, mode=\"all\")\n",
        "train_acc, val_acc, test_acc = test([graph_train, graph_val, graph_test], all_best_model)\n",
        "print('Best model:',\n",
        "      f'Train: {100 * train_acc:.2f}%, '\n",
        "      f'Valid: {100 * val_acc:.2f}% '\n",
        "      f'Test: {100 * test_acc:.2f}%')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JWkGiwB6Thr4"
      },
      "source": [
        "### 采样率为 0.8"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yWusJ9u3Tfhv"
      },
      "outputs": [],
      "source": [
        "args['ratios'] = (0.8, 0.8, 1)\n",
        "\n",
        "graphs_train, graphs_val, graphs_test = \\\n",
        "    GraphDataset.pyg_to_graphs(pyg_dataset, verbose=True, fixed_split=True)\n",
        "\n",
        "graph_train = graphs_train[0]\n",
        "graph_val = graphs_val[0]\n",
        "graph_test = graphs_test[0]\n",
        "\n",
        "model = GNN(graph_train.num_node_features, args['hidden_size'], graph_train.num_node_labels, args).to(args['device'])\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'])\n",
        "graphs = [graph_train, graph_val, graph_test]\n",
        "batch_best_model, batch_accs = train(graphs, graphs, args, model, optimizer)\n",
        "train_acc, val_acc, test_acc = test([graph_train, graph_val, graph_test], batch_best_model)\n",
        "print('Best model:',\n",
        "      f'Train: {100 * train_acc:.2f}%, '\n",
        "      f'Valid: {100 * val_acc:.2f}% '\n",
        "      f'Test: {100 * test_acc:.2f}%')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w_FjkNHDT4c6"
      },
      "source": [
        "### 采样率为 0.3"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "booJ6DASTjO4"
      },
      "outputs": [],
      "source": [
        "# Change the ratio to 0.3\n",
        "args['ratios'] = (0.3, 0.3, 1)\n",
        "\n",
        "graphs_train, graphs_val, graphs_test = \\\n",
        "    GraphDataset.pyg_to_graphs(pyg_dataset, verbose=True, fixed_split=True)\n",
        "\n",
        "graph_train = graphs_train[0]\n",
        "graph_val = graphs_val[0]\n",
        "graph_test = graphs_test[0]\n",
        "\n",
        "model = GNN(graph_train.num_node_features, args['hidden_size'], graph_train.num_node_labels, args).to(args['device'])\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'])\n",
        "graphs = [graph_train, graph_val, graph_test]\n",
        "batch_best_model, batch_accs_1 = train(graphs, graphs, args, model, optimizer)\n",
        "train_acc, val_acc, test_acc = test([graph_train, graph_val, graph_test], batch_best_model)\n",
        "print('Best model:',\n",
        "      f'Train: {100 * train_acc:.2f}%, '\n",
        "      f'Valid: {100 * val_acc:.2f}% '\n",
        "      f'Test: {100 * test_acc:.2f}%')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EePAvNlGUM2K"
      },
      "source": [
        "### 可视化\n",
        "\n",
        "下面所有的准确度都是在全批量模式下评估的。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7etNAkXAT55d"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "batch_results = np.array(batch_accs)\n",
        "batch_results_1 = np.array(batch_accs_1)\n",
        "all_results = np.array(all_accs)\n",
        "\n",
        "x = np.arange(1, 51)\n",
        "\n",
        "plt.figure(figsize=(9, 7))\n",
        "\n",
        "plt.plot(x, batch_results[:, 0], label=\"Batch 0.8 Train\")\n",
        "plt.plot(x, batch_results[:, 1], label=\"Batch 0.8 Validation\")\n",
        "plt.plot(x, batch_results[:, 2], label=\"Batch 0.8 Test\")\n",
        "plt.plot(x, batch_results_1[:, 0], label=\"Batch 0.3 Train\")\n",
        "plt.plot(x, batch_results_1[:, 1], label=\"Batch 0.3 Validation\")\n",
        "plt.plot(x, batch_results_1[:, 2], label=\"Batch 0.3 Test\")\n",
        "plt.plot(x, all_results[:, 0], label=\"All Train\")\n",
        "plt.plot(x, all_results[:, 1], label=\"All Validation\")\n",
        "plt.plot(x, all_results[:, 2], label=\"All Test\")\n",
        "plt.title('Model Accuracy')\n",
        "plt.ylabel('Accuracy')\n",
        "plt.xlabel('Epoch')\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iee0U8KGURc8"
      },
      "source": [
        "## 3 聚簇采样\n",
        "\n",
        "除了邻居采样，我们还可以用另一种方式来提升 GNN 的规模，即子图（聚簇）采样。这个方法是在 **Cluster-GCN**（[Chiang et al. (2019)](https://arxiv.org/abs/1905.07953)）中提出的。\n",
        "\n",
        "本小节中，我们会实现原始 **Cluster-GCN** 并在其上使用三个社群分割算法进行实验。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_BXjP79gUYir"
      },
      "source": [
        "### 配置"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UGQ_VKp8UOEm"
      },
      "outputs": [],
      "source": [
        "import copy\n",
        "import torch\n",
        "import random\n",
        "import numpy as np\n",
        "import networkx as nx\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import community as community_louvain\n",
        "\n",
        "from torch_geometric.nn import SAGEConv\n",
        "from torch.utils.data import DataLoader\n",
        "from torch_geometric.datasets import Planetoid\n",
        "from torch.nn import Sequential, Linear, ReLU\n",
        "from deepsnap.dataset import GraphDataset\n",
        "from deepsnap.graph import Graph\n",
        "\n",
        "pyg_dataset = Planetoid('./dataset', \"Cora\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bzMatyCSUaB6"
      },
      "outputs": [],
      "source": [
        "args = {\n",
        "    'device': torch.device('cuda' if torch.cuda.is_available() else 'cpu'),\n",
        "    'dropout': 0.5,\n",
        "    'num_layers': 2,\n",
        "    'hidden_size': 64,\n",
        "    'lr': 0.005,\n",
        "    'epochs': 150,\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ekV-sokSUeLc"
      },
      "source": [
        "### 将图分割成簇\n",
        "\n",
        "现在我们使用三个社群检测/分割算法来将图分成不同的簇：\n",
        "\n",
        "- [Kernighan–Lin algorithm (bisection)](https://networkx.org/documentation/stable/reference/algorithms/generated/networkx.algorithms.community.kernighan_lin.kernighan_lin_bisection.html)\n",
        "- [Clauset-Newman-Moore greedy modularity maximization](https://networkx.org/documentation/stable/reference/algorithms/generated/networkx.algorithms.community.modularity_max.greedy_modularity_communities.html#networkx.algorithms.community.modularity_max.greedy_modularity_communities)\n",
        "- [Louvain algorithm](https://python-louvain.readthedocs.io/en/latest/api.html)\n",
        "\n",
        "为了让训练更稳定，我们丢弃了结点少于 10 个的簇。\n",
        "\n",
        "现在将这些算法定义成 `DeepSNAP` 在图上的变换。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N8XeT005UcKh"
      },
      "outputs": [],
      "source": [
        "def preprocess(G, node_label_index, method=\"louvain\"):\n",
        "    graphs = []\n",
        "    labeled_nodes = set(node_label_index.tolist())\n",
        "    if method == \"louvain\":\n",
        "        community_mapping = community_louvain.best_partition(G, resolution=10)\n",
        "        communities = {}\n",
        "        for node in community_mapping:\n",
        "            comm = community_mapping[node]\n",
        "            if comm in communities:\n",
        "                communities[comm].add(node)\n",
        "            else:\n",
        "                communities[comm] = set([node])\n",
        "        communities = communities.values()\n",
        "    elif method == \"bisection\":\n",
        "        communities = nx.algorithms.community.kernighan_lin_bisection(G)\n",
        "    elif method == \"greedy\":\n",
        "        communities = nx.algorithms.community.greedy_modularity_communities(G)\n",
        "\n",
        "    for community in communities:\n",
        "        nodes = set(community)\n",
        "        subgraph = G.subgraph(nodes)\n",
        "        # Make sure each subgraph has more than 10 nodes\n",
        "        if subgraph.number_of_nodes() > 10:\n",
        "            node_mapping = {node : i for i, node in enumerate(subgraph.nodes())}\n",
        "            subgraph = nx.relabel_nodes(subgraph, node_mapping)\n",
        "            # Get the id of the training set labeled node in the new graph\n",
        "            train_label_index = []\n",
        "            for node in labeled_nodes:\n",
        "                if node in node_mapping:\n",
        "                    # Append relabeled labeled node index\n",
        "                    train_label_index.append(node_mapping[node])\n",
        "\n",
        "            # Make sure the subgraph contains at least one training set labeled node\n",
        "            if len(train_label_index) > 0:\n",
        "                dg = Graph(subgraph)\n",
        "                # Update node_label_index\n",
        "                dg.node_label_index = torch.tensor(train_label_index, dtype=torch.long)\n",
        "                graphs.append(dg)\n",
        "    return graphs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7CYEamCAU-TJ"
      },
      "source": [
        "### Louvain 处理"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-TrC6ybWU7eO"
      },
      "outputs": [],
      "source": [
        "graphs_train, graphs_val, graphs_test = \\\n",
        "    GraphDataset.pyg_to_graphs(pyg_dataset, verbose=True, fixed_split=True)\n",
        "\n",
        "graph_train = graphs_train[0]\n",
        "graph_val = graphs_val[0]\n",
        "graph_test = graphs_test[0]\n",
        "graphs = preprocess(graph_train.G, graph_train.node_label_index, method=\"louvain\")\n",
        "print(f\"Partition the graph in to {len(graphs)} communities\")\n",
        "avg_num_nodes = 0\n",
        "avg_num_edges = 0\n",
        "for graph in graphs:\n",
        "    avg_num_nodes += graph.num_nodes\n",
        "    avg_num_edges += graph.num_edges\n",
        "avg_num_nodes = int(avg_num_nodes / len(graphs))\n",
        "avg_num_edges = int(avg_num_edges / len(graphs))\n",
        "print(f\"Each community has {avg_num_nodes} nodes in average\")\n",
        "print(f\"Each community has {avg_num_edges} edges in average\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O03uXIuGVIgJ"
      },
      "source": [
        "### Louvain 训练"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iSbGf5ADVFQq"
      },
      "outputs": [],
      "source": [
        "model = GNN(graph_train.num_node_features, args['hidden_size'], graph_train.num_node_labels, args).to(args['device'])\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'])\n",
        "louvain_best_model, louvain_accs = train(graphs, [graph_train, graph_val, graph_test], args, model, optimizer, mode=\"community\")\n",
        "train_acc, val_acc, test_acc = test([graph_train, graph_val, graph_test], louvain_best_model)\n",
        "print('Best model:',\n",
        "      f'Train: {100 * train_acc:.2f}%, '\n",
        "      f'Valid: {100 * val_acc:.2f}% '\n",
        "      f'Test: {100 * test_acc:.2f}%')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6CvTf0ANVO9U"
      },
      "source": [
        "### Bisection 处理"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HkV0zlhgVJ7u"
      },
      "outputs": [],
      "source": [
        "graphs_train, graphs_val, graphs_test = \\\n",
        "    GraphDataset.pyg_to_graphs(pyg_dataset, verbose=True, fixed_split=True)\n",
        "\n",
        "graph_train = graphs_train[0]\n",
        "graph_val = graphs_val[0]\n",
        "graph_test = graphs_test[0]\n",
        "graphs = preprocess(graph_train.G, graph_train.node_label_index, method=\"bisection\")\n",
        "print(f\"Partition the graph in to {len(graphs)} communities\")\n",
        "avg_num_nodes = 0\n",
        "avg_num_edges = 0\n",
        "for graph in graphs:\n",
        "    avg_num_nodes += graph.num_nodes\n",
        "    avg_num_edges += graph.num_edges\n",
        "avg_num_nodes = int(avg_num_nodes / len(graphs))\n",
        "avg_num_edges = int(avg_num_edges / len(graphs))\n",
        "print(f\"Each community has {avg_num_nodes} nodes in average\")\n",
        "print(f\"Each community has {avg_num_edges} edges in average\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IqMCvP8wVVms"
      },
      "source": [
        "### Bisection 训练"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k1wgFg1bVRGY"
      },
      "outputs": [],
      "source": [
        "model = GNN(graph_train.num_node_features, args['hidden_size'], graph_train.num_node_labels, args).to(args['device'])\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'])\n",
        "bisection_best_model, bisection_accs = train(graphs, [graph_train, graph_val, graph_test], args, model, optimizer, mode=\"community\")\n",
        "train_acc, val_acc, test_acc = test([graph_train, graph_val, graph_test], bisection_best_model)\n",
        "print('Best model:',\n",
        "      f'Train: {100 * train_acc:.2f}%, '\n",
        "      f'Valid: {100 * val_acc:.2f}% '\n",
        "      f'Test: {100 * test_acc:.2f}%')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5PROPwoOVcJy"
      },
      "source": [
        "### Greedy 处理"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h3DVamWqVT92"
      },
      "outputs": [],
      "source": [
        "graphs_train, graphs_val, graphs_test = \\\n",
        "    GraphDataset.pyg_to_graphs(pyg_dataset, verbose=True, fixed_split=True)\n",
        "\n",
        "graph_train = graphs_train[0]\n",
        "graph_val = graphs_val[0]\n",
        "graph_test = graphs_test[0]\n",
        "graphs = preprocess(graph_train.G, graph_train.node_label_index, method=\"greedy\")\n",
        "print(f\"Partition the graph in to {len(graphs)} communities\")\n",
        "avg_num_nodes = 0\n",
        "avg_num_edges = 0\n",
        "for graph in graphs:\n",
        "    avg_num_nodes += graph.num_nodes\n",
        "    avg_num_edges += graph.num_edges\n",
        "avg_num_nodes = int(avg_num_nodes / len(graphs))\n",
        "avg_num_edges = int(avg_num_edges / len(graphs))\n",
        "print(f\"Each community has {avg_num_nodes} nodes in average\")\n",
        "print(f\"Each community has {avg_num_edges} edges in average\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "93pR_-kxVgma"
      },
      "source": [
        "### Greedy 训练"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lQgQY-jPVd_U"
      },
      "outputs": [],
      "source": [
        "model = GNN(graph_train.num_node_features, args['hidden_size'], graph_train.num_node_labels, args).to(args['device'])\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'])\n",
        "greedy_best_model, greedy_accs = train(graphs, [graph_train, graph_val, graph_test], args, model, optimizer, mode=\"community\")\n",
        "train_acc, val_acc, test_acc = test([graph_train, graph_val, graph_test], greedy_best_model)\n",
        "print('Best model:',\n",
        "      f'Train: {100 * train_acc:.2f}%, '\n",
        "      f'Valid: {100 * val_acc:.2f}% '\n",
        "      f'Test: {100 * test_acc:.2f}%')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D5edKKT6Vk1C"
      },
      "source": [
        "### 全批量训练"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N5tIXxC8ViFD"
      },
      "outputs": [],
      "source": [
        "graphs_train, graphs_val, graphs_test = \\\n",
        "    GraphDataset.pyg_to_graphs(pyg_dataset, verbose=True, fixed_split=True)\n",
        "\n",
        "graph_train = graphs_train[0]\n",
        "graph_val = graphs_val[0]\n",
        "graph_test = graphs_test[0]\n",
        "\n",
        "model = GNN(graph_train.num_node_features, args['hidden_size'], graph_train.num_node_labels, args).to(args['device'])\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=args['lr'])\n",
        "graphs = [graph_train, graph_val, graph_test]\n",
        "all_best_model, all_accs = train(graphs, graphs, args, model, optimizer, mode=\"all\")\n",
        "train_acc, val_acc, test_acc = test([graph_train, graph_val, graph_test], all_best_model)\n",
        "print('Best model:',\n",
        "      f'Train: {100 * train_acc:.2f}%, '\n",
        "      f'Valid: {100 * val_acc:.2f}% '\n",
        "      f'Test: {100 * test_acc:.2f}%')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6RpuETv7Vpx0"
      },
      "source": [
        "### 可视化"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PMK33kY5VmF5"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "louvain_results = np.array(louvain_accs)\n",
        "bisection_results = np.array(bisection_accs)\n",
        "greedy_results = np.array(greedy_accs)\n",
        "all_results = np.array(all_accs)\n",
        "\n",
        "x = np.arange(1, 151)\n",
        "\n",
        "plt.figure(figsize=(9, 7))\n",
        "\n",
        "plt.plot(x, louvain_results[:, 1], label=\"Louvain Validation\")\n",
        "plt.plot(x, louvain_results[:, 2], label=\"Louvain Test\")\n",
        "plt.plot(x, bisection_results[:, 1], label=\"Bisection Validation\")\n",
        "plt.plot(x, bisection_results[:, 2], label=\"Bisection Test\")\n",
        "plt.plot(x, greedy_results[:, 1], label=\"Greedy Validation\")\n",
        "plt.plot(x, greedy_results[:, 2], label=\"Greedy Test\")\n",
        "plt.plot(x, all_results[:, 1], label=\"All Validation\")\n",
        "plt.plot(x, all_results[:, 2], label=\"All Test\")\n",
        "plt.title('Model Accuracy')\n",
        "plt.ylabel('Accuracy')\n",
        "plt.xlabel('Epoch')\n",
        "plt.legend()\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "CS224W - Colab 5.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3.9.12 ('cs224w')",
      "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.9.12"
    },
    "vscode": {
      "interpreter": {
        "hash": "3d401f214986802b882e8a583f41ee2435f6c7e7677ce5b0132581f351a1fd17"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}