tutorials/01_node_classification_on_citation_network.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Node Classification on Citation Network\n",
    "\n",
    "As a start, we present a end to end example, demonstrating how GraphScope process node classification task on citation network by combining analytics, interactive and graph neural networks computation.\n",
    "\n",
    "In this example, we use [ogbn-mag](https://ogb.stanford.edu/docs/nodeprop/#ogbn-mag) dataset. ogbn-mag is a heterogeneous network composed of a subset of the Microsoft Academic Graph. It contains 4 types of entities(i.e., papers, authors, institutions, and fields of study), as well as four types of directed relations connecting two entities.\n",
    "\n",
    "Given the heterogeneous ogbn-mag data, the task is to predict the class of each paper. We apply both the attribute and structural information to classify papers. In the graph, each paper node contains a 128-dimensional word2vec vector representing its content, which is obtained by averaging the embeddings of words in its title and abstract. The embeddings of individual words are pre-trained. The structural information is computed on-the-fly.\n",
    "\n",
    "This tutorial has the following steps:\n",
    "- Querying graph data with Gremlin;\n",
    "- Running graph analytical algorithms;\n",
    "- Running graph-based machine learning tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install graphscope package if you are NOT in the Playground\n",
    "\n",
    "!pip3 install graphscope\n",
    "!pip3 uninstall -y importlib_metadata  # Address an module conflict issue on colab.google. Remove this line if you are not on colab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the graphscope module\n",
    "\n",
    "import graphscope\n",
    "\n",
    "graphscope.set_option(show_log=False)  # enable logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the obgn_mag dataset as a graph\n",
    "\n",
    "from graphscope.dataset import load_ogbn_mag\n",
    "\n",
    "graph = load_ogbn_mag()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interactive query with gremlin\n",
    "\n",
    "In this example, we launch a interactive query and use graph traversal to count the number of papers two given authors have co-authored. To simplify the query, we assume the authors can be uniquely identified by ID `2` and `4307`, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the entrypoint for submitting Gremlin queries on graph g.\n",
    "interactive = graphscope.gremlin(graph)\n",
    "\n",
    "# Count the number of papers two authors (with id 2 and 4307) have co-authored.\n",
    "papers = interactive.execute(\n",
    "    \"g.V().has('author', 'id', 2).out('writes').where(__.in('writes').has('id', 4307)).count()\"\n",
    ").one()\n",
    "print(\"result\", papers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph analytics with analytical engine\n",
    "\n",
    "Continuing our example, we run graph algorithms on graph to generate structural features. below we first derive a subgraph by extracting publications in specific time out of the entire graph (using Gremlin!), and then run k-core decomposition and triangle counting to generate the structural features of each paper node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exact a subgraph of publication within a time range.\n",
    "sub_graph = interactive.subgraph(\"g.V().has('year', inside(2014, 2020)).outE('cites')\")\n",
    "\n",
    "# Project the subgraph to simple graph by selecting papers and their citations.\n",
    "simple_g = sub_graph.project(vertices={\"paper\": []}, edges={\"cites\": []})\n",
    "# compute the kcore and triangle-counting.\n",
    "kc_result = graphscope.k_core(simple_g, k=5)\n",
    "tc_result = graphscope.triangles(simple_g)\n",
    "\n",
    "# Add the results as new columns to the citation graph.\n",
    "sub_graph = sub_graph.add_column(kc_result, {\"kcore\": \"r\"})\n",
    "sub_graph = sub_graph.add_column(tc_result, {\"tc\": \"r\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph neural networks (GNNs)\n",
    "\n",
    "Then, we use the generated structural features and original features to train a learning model with learning engine.\n",
    "\n",
    "In our example, we train a supervised GraphSAGE model to classify the nodes (papers) into 349 categories,\n",
    "each of which represents a venue (e.g. pre-print and conference)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the features for learning,\n",
    "# we chose original 128-dimension feature and k-core, triangle count result as new features.\n",
    "paper_features = []\n",
    "for i in range(128):\n",
    "    paper_features.append(\"feat_\" + str(i))\n",
    "paper_features.append(\"kcore\")\n",
    "paper_features.append(\"tc\")\n",
    "\n",
    "# Launch a learning engine. here we split the dataset, 75% as train, 10% as validation and 15% as test.\n",
    "lg = graphscope.graphlearn(\n",
    "    sub_graph,\n",
    "    nodes=[(\"paper\", paper_features)],\n",
    "    edges=[(\"paper\", \"cites\", \"paper\")],\n",
    "    gen_labels=[\n",
    "        (\"train\", \"paper\", 100, (0, 75)),\n",
    "        (\"val\", \"paper\", 100, (75, 85)),\n",
    "        (\"test\", \"paper\", 100, (85, 100)),\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Then we define the training process using the example GraphSAGE model with tensorflow.\n",
    "try:\n",
    "    # https://www.tensorflow.org/guide/migrate\n",
    "    import tensorflow.compat.v1 as tf\n",
    "    tf.disable_v2_behavior()\n",
    "except ImportError:\n",
    "    import tensorflow as tf\n",
    "\n",
    "import graphscope.learning\n",
    "from graphscope.learning.examples import EgoGraphSAGE\n",
    "from graphscope.learning.examples import EgoSAGESupervisedDataLoader\n",
    "from graphscope.learning.examples.tf.trainer import LocalTrainer\n",
    "\n",
    "# supervised SAGE.\n",
    "def train_sage(graph, node_type, edge_type, class_num, features_num,\n",
    "              hops_num=2, nbrs_num=[25, 10], epochs=2,\n",
    "              hidden_dim=256, in_drop_rate=0.5, learning_rate=0.01,\n",
    "):\n",
    "    graphscope.learning.reset_default_tf_graph()\n",
    "\n",
    "    dimensions = [features_num] + [hidden_dim] * (hops_num - 1) + [class_num]\n",
    "    model = EgoGraphSAGE(dimensions, act_func=tf.nn.relu, dropout=in_drop_rate)\n",
    "\n",
    "    # prepare train dataset\n",
    "    train_data = EgoSAGESupervisedDataLoader(\n",
    "        graph, graphscope.learning.Mask.TRAIN,\n",
    "        node_type=node_type, edge_type=edge_type, nbrs_num=nbrs_num, hops_num=hops_num,\n",
    "    )\n",
    "    train_embedding = model.forward(train_data.src_ego)\n",
    "    train_labels = train_data.src_ego.src.labels\n",
    "    loss = tf.reduce_mean(\n",
    "        tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
    "            labels=train_labels, logits=train_embedding,\n",
    "        )\n",
    "    )\n",
    "    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)\n",
    "\n",
    "    # prepare test dataset\n",
    "    test_data = EgoSAGESupervisedDataLoader(\n",
    "        graph, graphscope.learning.Mask.TEST,\n",
    "        node_type=node_type, edge_type=edge_type, nbrs_num=nbrs_num, hops_num=hops_num,\n",
    "    )\n",
    "    test_embedding = model.forward(test_data.src_ego)\n",
    "    test_labels = test_data.src_ego.src.labels\n",
    "    test_indices = tf.math.argmax(test_embedding, 1, output_type=tf.int32)\n",
    "    test_acc = tf.div(\n",
    "        tf.reduce_sum(tf.cast(tf.math.equal(test_indices, test_labels), tf.float32)),\n",
    "        tf.cast(tf.shape(test_labels)[0], tf.float32),\n",
    "    )\n",
    "\n",
    "    # train and test\n",
    "    trainer = LocalTrainer()\n",
    "    trainer.train(train_data.iterator, loss, optimizer, epochs=epochs)\n",
    "    trainer.test(test_data.iterator, test_acc)\n",
    "\n",
    "train_sage(lg, node_type=\"paper\", edge_type=\"cites\",\n",
    "          class_num=349,  # output dimension\n",
    "          features_num=130,  # input dimension, 128 + kcore + triangle count\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}