tutorials/08_unsupervised_learning_with_graphsage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unsupervised Graph Learning with GraphSage\n",
    "\n",
    "\n",
    "GraphScope provides the capability to process learning tasks. In this tutorial, we demonstrate how GraphScope trains a model with GraphSage.\n",
    "\n",
    "The task is link prediction, which estimates the probability of links between nodes in a graph.\n",
    "\n",
    "In this task, we use our implementation of GraphSAGE algorithm to build a model that predicts protein-protein links in the [PPI](https://humgenomics.biomedcentral.com/articles/10.1186/1479-7364-3-3-291) dataset. In which every node represents a protein. The task can be treated as a unsupervised link prediction on a homogeneous link network.\n",
    "\n",
    "In this task, GraphSage algorithm would compress both structural and attribute information in the graph into low-dimensional embedding vectors on each node. These embeddings can be further used to predict links between nodes.\n",
    "\n",
    "This tutorial has following steps:\n",
    "- Launching the learning engine and attaching to loaded graph.\n",
    "- Defining train process with builtin GraphSage model and hyper-parameters\n",
    "- Training and evaluating\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install graphscope package if you are NOT in the Playground\n",
    "\n",
    "!pip3 install graphscope"
   ]
  },
  {
   "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 ppi dataset\n",
    "\n",
    "from graphscope.dataset import load_ppi\n",
    "\n",
    "graph = load_ppi()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launch learning engine \n",
    "Then, we need to define a feature list for training. The training feature list should be selected from the vertex properties. In this case, we choose all the properties prefix with \"feat-\" as the training features.\n",
    "\n",
    "With the feature list, next we launch a learning engine with the [graphlearn](https://graphscope.io/docs/reference/session.html#graphscope.Session.graphlearn) method of graphscope.\n",
    "\n",
    "In this case, we specify the GraphSAGE training over \"protein\" nodes and \"link\" edges.\n",
    "\n",
    "With gen_labels, we take all the protein nodes as training set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the features for learning\n",
    "protein_features = []\n",
    "for i in range(50):\n",
    "    protein_features.append(\"feat-\" + str(i))\n",
    "\n",
    "# launch a learning engine.\n",
    "lg = graphscope.graphlearn(\n",
    "    graph,\n",
    "    nodes=[(\"protein\", protein_features)],\n",
    "    edges=[(\"protein\", \"link\", \"protein\")],\n",
    "    gen_labels=[\n",
    "        (\"train\", \"protein\", 100, (0, 100)),\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "We use the builtin GraphSAGE model to define the training process.\n",
    "\n",
    "In the example, we use tensorflow as \"NN\" backend trainer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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 argparse\n",
    "import graphscope.learning.graphlearn.python.nn.tf as tfg\n",
    "from graphscope.learning.examples import EgoGraphSAGE\n",
    "from graphscope.learning.examples import EgoSAGEUnsupervisedDataLoader\n",
    "from graphscope.learning.examples.tf.trainer import LocalTrainer\n",
    "\n",
    "def parse_args():\n",
    "  argparser = argparse.ArgumentParser(\"Train EgoSAGE Unsupervised.\")\n",
    "  argparser.add_argument('--batch_size', type=int, default=512)\n",
    "  argparser.add_argument('--features_num', type=int, default=50)\n",
    "  argparser.add_argument('--hidden_dim', type=int, default=128)\n",
    "  argparser.add_argument('--output_dim', type=int, default=128)\n",
    "  argparser.add_argument('--nbrs_num', type=list, default=[5, 5])\n",
    "  argparser.add_argument('--learning_rate', type=float, default=0.01)\n",
    "  argparser.add_argument('--epochs', type=int, default=2)\n",
    "  argparser.add_argument('--drop_out', type=float, default=0.0)\n",
    "  argparser.add_argument('--temperature', type=float, default=0.07)\n",
    "  argparser.add_argument('--node_type', type=str, default=\"protein\")\n",
    "  argparser.add_argument('--edge_type', type=str, default=\"link\")\n",
    "  return argparser.parse_args()\n",
    "args = parse_args()\n",
    "\n",
    "# Define Model\n",
    "dims = [args.features_num] + [args.hidden_dim] * (len(args.nbrs_num) - 1) + [args.output_dim]\n",
    "model = EgoGraphSAGE(dims, dropout=args.drop_out)\n",
    "\n",
    "# Prepare train dataset\n",
    "train_data = EgoSAGEUnsupervisedDataLoader(lg, None, batch_size=args.batch_size,\n",
    "    node_type=args.node_type, edge_type=args.edge_type, nbrs_num=args.nbrs_num)\n",
    "src_emb = model.forward(train_data.src_ego)\n",
    "dst_emb = model.forward(train_data.dst_ego)\n",
    "neg_dst_emb = model.forward(train_data.neg_dst_ego)\n",
    "loss = tfg.unsupervised_softmax_cross_entropy_loss(\n",
    "    src_emb, dst_emb, neg_dst_emb, temperature=args.temperature)\n",
    "optimizer = tf.train.AdamOptimizer(learning_rate=args.learning_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run training process\n",
    "\n",
    "After define training process and hyperparameters, we can start training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = LocalTrainer()\n",
    "trainer.train(train_data.iterator, loss, optimizer, epochs=args.epochs)"
   ]
  }
 ],
 "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
}