12.2.2 GraphSAGE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. 超参数设置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "lr = 0.001\n",
    "epochs = 20\n",
    "hidden_dim = 128\n",
    "dataset_name = 'cora-seeds'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. 数据集处理与加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = gb.BuiltinDataset(dataset_name).load()\n",
    "graph = dataset.graph.to(device)\n",
    "feature = dataset.feature.to(device)\n",
    "train_set = dataset.tasks[1].train_set\n",
    "test_set = dataset.tasks[1].test_set\n",
    "task_name = dataset.tasks[1].metadata[\"name\"]\n",
    " \n",
    "datapipe = gb.ItemSampler(train_set, batch_size=256, shuffle=True)\n",
    "datapipe = datapipe.copy_to(device)\n",
    "datapipe = datapipe.sample_uniform_negative(graph, 5)\n",
    "datapipe = datapipe.sample_neighbor(graph, [5, 5])\n",
    "datapipe = datapipe.transform(partial(gb.exclude_seed_edges, include_reverse_edges=True))\n",
    "datapipe = datapipe.fetch_feature(feature, node_feature_keys=[\"feat\"])\n",
    "datapipe = datapipe.copy_to(device)\n",
    "dataloader = gb.DataLoader(datapipe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. 构建卷积层和模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAGEConv(nn.Module):\n",
    "    def __init__(self, in_feats, out_feats):\n",
    "        super(SAGEConv, self).__init__()\n",
    "        # 定义GraphSAGE参数\n",
    "        self.fc_neigh = nn.Linear(in_feats, out_feats)\n",
    "        self.fc_self = nn.Linear(in_feats, out_feats)\n",
    "        # 参数初始化\n",
    "        gain = nn.init.calculate_gain(\"relu\")\n",
    "        nn.init.xavier_uniform_(self.fc_neigh.weight, gain=gain)\n",
    "        nn.init.xavier_uniform_(self.fc_self.weight, gain=gain)\n",
    "    def forward(self, graph, feat):\n",
    "        with graph.local_scope():\n",
    "            # 获取源节点和目标节点的特征\n",
    "            feat_src = feat_dst = feat\n",
    "            feat_dst = feat_src[: graph.number_of_dst_nodes()]\n",
    "            # 定义消息传递函数\n",
    "            msg_fn = fn.copy_u(\"h\", \"m\")\n",
    "            # 定义聚合函数\n",
    "            reduce_fn = fn.mean(\"m\", \"neigh\")\n",
    "            # 自身特征\n",
    "            h_self = feat_dst\n",
    "            # 消息传递\n",
    "            graph.srcdata[\"h\"] = self.fc_neigh(feat_src)\n",
    "            graph.update_all(msg_fn, reduce_fn)\n",
    "            h_neigh = graph.dstdata[\"neigh\"]\n",
    "            # 聚合自身特征和邻居特征\n",
    "            rst = self.fc_self(h_self) + h_neigh\n",
    "            return rst\n",
    "\n",
    "# 定义 GraphSAGE 模型\n",
    "class SAGE(nn.Module):\n",
    "    def __init__(self, in_size, hidden_size):\n",
    "        super().__init__()\n",
    "        self.layers = nn.ModuleList()\n",
    "        # 增加graphsage层\n",
    "        self.layers.append(SAGEConv(in_size, hidden_size))\n",
    "        self.layers.append(SAGEConv(hidden_size, hidden_size))\n",
    "        self.layers.append(SAGEConv(hidden_size, hidden_size))\n",
    "        self.hidden_size = hidden_size\n",
    "        # 定义预测器\n",
    "        self.predictor = nn.Sequential(\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, 1),\n",
    "        )\n",
    "    def forward(self, blocks, x):\n",
    "        # 前向传播获得节点嵌入\n",
    "        hidden_x = x\n",
    "        for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)):\n",
    "            hidden_x = layer(block, hidden_x)\n",
    "            is_last_layer = layer_idx == len(self.layers) - 1\n",
    "            if not is_last_layer:\n",
    "                hidden_x = F.relu(hidden_x)\n",
    "        return hidden_x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. 定义损失函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.BCEWithLogitsLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. 设置优化器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=lr)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6. 模型评估函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, dataloader):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for step, data in enumerate(dataloader):\n",
    "            compacted_seeds = data.compacted_seeds.T\n",
    "            labels = data.labels\n",
    "            node_feature = data.node_features[\"feat\"]\n",
    "            blocks = data.blocks\n",
    "            # 获得节点embedding\n",
    "            y = model(blocks, node_feature)\n",
    "            # 使用预测器获得预测值\n",
    "            logits = model.predictor(\n",
    "                y[compacted_seeds[0]] * y[compacted_seeds[1]]\n",
    "            ).squeeze()\n",
    "            # 获得预测结果\n",
    "            preds = torch.round(torch.sigmoid(logits))\n",
    "            # 预测准确的个数\n",
    "            correct += (preds == labels).sum().item()\n",
    "            # 样本总数\n",
    "            total += labels.size(0)\n",
    "    # 计算准确率\n",
    "    accuracy = correct / total\n",
    "    return accuracy\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7. 模型训练流程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for epoch in range(epochs):\n",
    "    model.train()\n",
    "    total_loss = 0\n",
    "    start_epoch_time = time.time()\n",
    "    for step, data in enumerate(dataloader):\n",
    "        # mini-batch中的节点ID\n",
    "        compacted_seeds = data.compacted_seeds.T\n",
    "        # 获取mini-batch中的标签、特征、block\n",
    "        labels = data.labels\n",
    "        node_feature = data.node_features[\"feat\"]\n",
    "        blocks = data.blocks\n",
    "        # 获取输入节点的表示\n",
    "        y = model(blocks, node_feature)\n",
    "        logits = model.predictor(\n",
    "            y[compacted_seeds[0]] * y[compacted_seeds[1]]\n",
    "        ).squeeze()\n",
    "        # 计算损失并反向传播\n",
    "        loss = criterion(logits, labels)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        total_loss += loss.item()\n",
    "    end_epoch_time = time.time()\n",
    "    print(f\"Epoch {epoch} | Loss: {total_loss:.4f} | Time: {end_epoch_time - start_epoch_time:.2f}s\")\n",
    "    # 评估训练结果\n",
    "    if (epoch+1) % 5 == 0:\n",
    "        accuracy = evaluate(model, dataloader)\n",
    "        print(f\"Accuracy: {accuracy:.4f}\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ggl2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}