{ "cells": [ { "attachments": {}, "cell_type": "markdown", "id": "2a09f4f5-e147-4024-9deb-1e37753d14f9", "metadata": {}, "source": [ "Copyright (c) 2023 Graphcore Ltd. All rights reserved." ] }, { "attachments": {}, "cell_type": "markdown", "id": "73b4dca6-3257-4709-aac0-9ae6441216aa", "metadata": {}, "source": [ "# Training Neural Bellman-Ford Networks (NBFnet) for Inductive Knowledge Graph Link Prediction on IPUs \n", "\n", "Neural Bellman-Ford networks (NBFNet) is a model that generalises path-based reasoning models for predicting links in homogeneous and heterogeneous graphs. \n", "\n", "In this notebook we use NBFNet for link prediction in the FB15k-237 knowledge graph with 14541 entities, 237 relation types and 272115 triples. However in practice we explicitly insert reverse edges, which brings us to a total of 474 relation types and 544230 triples.\n", "\n", "Unlike many other knowledge graph completion models, NBFNet can be *inductive*, in other words it can generalise to entities that do not appear in the training data. To demonstrate this inductive behaviour we train the model on a small subset of the graph (4707 entities, 54406 triples) and perform inference on the complete FB15k-237 graph.\n", "\n", "| Domain | Tasks | Model | Datasets | Workflow | Number of IPUs | Execution time |\n", "|---------|-------|-------|----------|----------|--------------------|----------------|\n", "| GNNs | Link Prediction | NBFNet | FB15k-237 | Training, evaluation | recommended: 16 (min: 4) | 20mn |\n", "\n", "This notebook assumes some familiarity with PopTorch as well as PyTorch Geometric (PyG). For additional resources please consult:\n", "* [PopTorch Documentation](https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/index.html),\n", "* [PopTorch Examples and Tutorials](https://docs.graphcore.ai/en/latest/examples.html#pytorch),\n", "* [PyTorch Geometric](https://pytorch-geometric.readthedocs.io/en/latest/),\n", "* [PopTorch Geometric Documentation](https://docs.graphcore.ai/projects/poptorch-geometric-user-guide/en/latest/index.html),\n", "\n", "## Environment setup\n", "\n", "The best way to run this demo is on Paperspace Gradient’s cloud IPUs because everything is already set up for you. To improve your experience, we preload datasets and pre-install packages. This can take a few minutes. If you experience errors immediately after starting a session, please try restarting the kernel before contacting support. If a problem persists or you want to give us feedback on the content of this notebook, please reach out to through our community of developers using our [Slack channel](https://www.graphcore.ai/join-community) or raise a [GitHub issue](https://github.com/graphcore/Gradient-PyTorch-Geometric/issues).\n", "\n", "\n", "To run the demo using other IPU hardware, you need to have the Poplar SDK enabled. Refer to the [Getting Started guide](https://docs.graphcore.ai/en/latest/getting-started.html#getting-started) for your system for details on how to enable the Poplar SDK. Also refer to the [Jupyter Quick Start guide](https://docs.graphcore.ai/projects/jupyter-notebook-quick-start/en/latest/index.html) for how to set up Jupyter to be able to run this notebook on a remote IPU machine.\n", "\n", "Requirements:\n", "\n", "* Python packages installed with `pip install -r requirements.txt`" ] }, { "attachments": {}, "cell_type": "markdown", "id": "9d97680e", "metadata": {}, "source": [ "In order to improve usability and support for future users, Graphcore would like to collect information about the\n", "applications and code being run in this notebook. The following information will be anonymised before being sent to Graphcore:\n", "\n", "- User progression through the notebook\n", "- Notebook details: number of cells, code being run and the output of the cells\n", "- Environment details\n", "\n", "You can disable logging at any time by running `%unload_ext graphcore_cloud_tools.notebook_logging.gc_logger` from any cell." ] }, { "cell_type": "code", "execution_count": null, "id": "622cd22c-178b-48df-a6b8-6a17cea06cd3", "metadata": {}, "outputs": [], "source": [ "%pip install -r requirements.txt\n", "%load_ext graphcore_cloud_tools.notebook_logging.gc_logger" ] }, { "attachments": {}, "cell_type": "markdown", "id": "4b89d3b9", "metadata": {}, "source": [ "Let's import the required packages:" ] }, { "cell_type": "code", "execution_count": null, "id": "c10605e9-a1af-4a6b-affd-7f2b046c9ca7", "metadata": {}, "outputs": [], "source": [ "import os\n", "import os.path as osp\n", "\n", "import poptorch\n", "import torch\n", "from torch_geometric.datasets import RelLinkPredDataset\n", "\n", "import data as nbfnet_data\n", "import inference_utils\n", "from nbfnet import NBFNet" ] }, { "attachments": {}, "cell_type": "markdown", "id": "7d840e00-805e-4efe-b6a9-c1c16e6fbb28", "metadata": {}, "source": [ "For compatibility with the Paperspace environment variables we need to do the following:" ] }, { "cell_type": "code", "execution_count": null, "id": "9040c772-cb87-4b81-8e76-bbd429209bb9", "metadata": {}, "outputs": [], "source": [ "poptorch.setLogLevel(\"ERR\")\n", "executable_cache_dir = (\n", " os.getenv(\"POPLAR_EXECUTABLE_CACHE_DIR\", \"/tmp/exe_cache/\") + \"/pyg-nbfnet\"\n", ")\n", "dataset_directory = os.getenv(\"DATASETS_DIR\", \"data\")\n", "available_ipus = int(os.getenv(\"NUM_AVAILABLE_IPU\", \"4\"))" ] }, { "attachments": {}, "cell_type": "markdown", "id": "6be914e8", "metadata": {}, "source": [ "Now we are ready to start!" ] }, { "attachments": {}, "cell_type": "markdown", "id": "faa8d5c6-9ef5-4625-86c7-01645abe4b42", "metadata": {}, "source": [ "## Training the model\n", "\n", "First we will look at the steps required to train the model.\n", "\n", "### Defining the model hyperparameters\n", "\n", "Here we define some model settings and hyperparameters:\n", "- BATCH_SIZE: The micro batch size (number of triples `(head, relation, tail)`) during training\n", "- NUM_NEGATIVES: The number of triples `(head, relation, false_tail)` to contrast against each true triple\n", "- LEARNING_RATE\n", "- LATENT_DIM: The hidden dimension in the Message Passing Neural Network\n", "- NUM_LAYERS: The number of message passing layers\n", "- NEG_ADVERSARIAL_TEMP: The temperature of a softmax that weights negative samples based on their difficulty." ] }, { "cell_type": "code", "execution_count": null, "id": "27d2cf4e-fc68-41c0-8fee-bb2cfeaf5adf", "metadata": {}, "outputs": [], "source": [ "BATCH_SIZE = 6\n", "NUM_NEGATIVES = 32\n", "LEARNING_RATE = 0.001\n", "LATENT_DIM = 64\n", "NUM_LAYERS = 6\n", "NEG_ADVERSARIAL_TEMP = 0.7" ] }, { "attachments": {}, "cell_type": "markdown", "id": "43317bab-c1cc-48c8-9228-afea2b8f50e1", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "### Creating the dataset and dataloader\n", "\n", "Now we build a training and validation dataset from the small `IndFB15k-237_v4` graph and a test dataset from the full `FB15k-237` graph. Then we create a dataloader for training, validation and test. The dataloader does the following: batches data; removes edges between head and tail entities in the training dataset to make the training objective non-trivial; samples negative tails. For validation and test, all entities will be treated as potential tail nodes." ] }, { "cell_type": "code", "execution_count": null, "id": "35952997-6009-4d59-a13f-c0fef7dc60a3", "metadata": {}, "outputs": [], "source": [ "dataset_train = nbfnet_data.build_dataset(\n", " name=\"IndFB15k-237\", path=dataset_directory, version=\"v4\"\n", ")\n", "dataset_inference = nbfnet_data.build_dataset(name=\"FB15k-237\", path=dataset_directory)" ] }, { "cell_type": "code", "execution_count": null, "id": "357c4565-43ce-43ed-b986-4c5443d714ce", "metadata": {}, "outputs": [], "source": [ "dataloader = dict(\n", " train=nbfnet_data.DataWrapper(\n", " nbfnet_data.NBFData(\n", " data=dataset_train[0],\n", " batch_size=BATCH_SIZE,\n", " is_training=True,\n", " num_relations=dataset_train.num_relations,\n", " num_negatives=NUM_NEGATIVES,\n", " )\n", " ),\n", " valid=nbfnet_data.DataWrapper(\n", " nbfnet_data.NBFData(\n", " data=dataset_train[1],\n", " batch_size=1,\n", " is_training=False,\n", " )\n", " ),\n", " test=nbfnet_data.DataWrapper(\n", " nbfnet_data.NBFData(\n", " data=dataset_inference[2],\n", " batch_size=1,\n", " is_training=False,\n", " )\n", " ),\n", ")\n", "\n", "num_relations = dataset_inference.num_relations + 1" ] }, { "attachments": {}, "cell_type": "markdown", "id": "e69d12aa-9444-4bec-b971-cc350dcd01f6", "metadata": {}, "source": [ "### Defining the model\n", "We can now define the model and the optimiser using the hyperparameters that we have defined above. The model is cast to float16 for improved compute- and memory efficiency." ] }, { "cell_type": "code", "execution_count": null, "id": "db1c2480-9a8c-4d43-94e7-a2e5f7f6fc73", "metadata": {}, "outputs": [], "source": [ "model = NBFNet(\n", " input_dim=LATENT_DIM,\n", " hidden_dims=[LATENT_DIM] * NUM_LAYERS,\n", " message_fct=\"mult\",\n", " aggregation_fct=\"sum\",\n", " num_mlp_layers=2,\n", " relation_learning=\"linear_query\",\n", " adversarial_temperature=NEG_ADVERSARIAL_TEMP,\n", " num_relations=num_relations,\n", ")\n", "\n", "model.half();" ] }, { "cell_type": "code", "execution_count": null, "id": "12cfd3a0-d41b-4666-81d0-9ef76a586c27", "metadata": {}, "outputs": [], "source": [ "optim = poptorch.optim.AdamW(\n", " model.parameters(),\n", " lr=LEARNING_RATE,\n", " bias_correction=True,\n", " weight_decay=0.0,\n", " eps=1e-8,\n", " betas=(0.9, 0.999),\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "a2c5dc76-2100-4f3a-bc4c-f1eb3dfc1344", "metadata": { "tags": [] }, "source": [ "The model defines a `poptorch.Stage` for every layer as well as for the preprocessing and prediction step. We can now define the IPUs to every block for a pipelined (or sharded) execution:" ] }, { "cell_type": "code", "execution_count": null, "id": "1a61cc60-dc46-4bd5-b106-cdf38f26eae7", "metadata": {}, "outputs": [], "source": [ "pipeline = {\n", " \"preprocessing\": 0,\n", " \"layer0\": 0,\n", " \"layer1\": 1,\n", " \"layer2\": 1,\n", " \"layer3\": 2,\n", " \"layer4\": 2,\n", " \"layer5\": 3,\n", " \"prediction\": 3,\n", "}" ] }, { "attachments": {}, "cell_type": "markdown", "id": "7891a1b3", "metadata": {}, "source": [ "And assign them using the PopTorch options:" ] }, { "cell_type": "code", "execution_count": null, "id": "c53911d6-f3a0-4ee5-a70a-4eea85c43636", "metadata": {}, "outputs": [], "source": [ "pipeline_plan = [poptorch.Stage(k).ipu(v) for k, v in pipeline.items()]\n", "\n", "train_opts = poptorch.Options()\n", "train_opts.setExecutionStrategy(poptorch.PipelinedExecution(*pipeline_plan))\n", "train_opts.Training.gradientAccumulation(16 if available_ipus == 16 else 64)\n", "if available_ipus == 16:\n", " train_opts.replicationFactor(4)\n", "train_opts.enableExecutableCaching(executable_cache_dir)\n", "\n", "test_opts = poptorch.Options()\n", "test_opts.setExecutionStrategy(poptorch.PipelinedExecution(*pipeline_plan))\n", "test_opts.deviceIterations(len(set(pipeline.values())))\n", "test_opts.enableExecutableCaching(executable_cache_dir)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "dd8fe613-f4f4-4458-b9e0-a6d8e4985499", "metadata": {}, "source": [ "We wrap the dataloader into a `poptorch.DataLoader`:" ] }, { "cell_type": "code", "execution_count": null, "id": "ec3f607c-5c1e-428e-8126-24e6080916b1", "metadata": {}, "outputs": [], "source": [ "for partition in [\"train\", \"valid\", \"test\"]:\n", " dataloader[partition] = poptorch.DataLoader(\n", " options=train_opts if partition == \"train\" else test_opts,\n", " dataset=dataloader[partition],\n", " batch_size=1,\n", " collate_fn=nbfnet_data.custom_collate,\n", " )" ] }, { "attachments": {}, "cell_type": "markdown", "id": "e0abbed9", "metadata": {}, "source": [ "And wrap the model into `poptorch.trainingModel` or `poptorch.inferenceModel`:" ] }, { "cell_type": "code", "execution_count": null, "id": "3cdd7991-00de-4e6d-8963-2154c3edc6c1", "metadata": {}, "outputs": [], "source": [ "model_train = poptorch.trainingModel(model, options=train_opts, optimizer=optim)\n", "model_valid = poptorch.inferenceModel(model, options=test_opts)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "f27be028", "metadata": {}, "source": [ "Now we are ready to start training." ] }, { "attachments": {}, "cell_type": "markdown", "id": "8b44cac1-4b3d-4a43-8ad9-f71b046cf76f", "metadata": {}, "source": [ "### Training the model\n", "Now we are ready to train the model. We run training for 5 epochs on the IndFB15k-237_v4 subgraph with interleaved validation." ] }, { "cell_type": "code", "execution_count": null, "id": "735d8e1f", "metadata": {}, "outputs": [], "source": [ "num_epochs = 5" ] }, { "cell_type": "code", "execution_count": null, "id": "f8172c5c-6bbe-401c-962f-bc9256f5e9b8", "metadata": {}, "outputs": [], "source": [ "model_train.train()\n", "model_valid.eval()\n", "\n", "loss_per_epoch = []\n", "mrr_per_epoch = []\n", "for epoch in range(num_epochs):\n", " total_loss = 0\n", " total_count = 0\n", " for batch in dataloader[\"train\"]:\n", " loss, count = model_train(**batch)\n", " loss, count = loss.mean(), count.sum() # reduction across replicas\n", " total_loss += float(loss) * count\n", " total_count += count\n", " loss_per_epoch.append(total_loss / total_count)\n", " print(f\"Epoch {epoch} finished, training loss {total_loss / total_count:.4}\")\n", "\n", " # Interleaved validation\n", " mrr = 0\n", " total_count = 0\n", " model_train.detachFromDevice()\n", " for batch in dataloader[\"valid\"]:\n", " prediction, count, mask, _ = model_valid(**batch)\n", " if isinstance(count, torch.Tensor):\n", " count = count.sum()\n", " prediction = prediction[mask]\n", " true_score = prediction[:, 0:1]\n", " rank = torch.sum(true_score <= prediction, dim=-1)\n", " mrr += float(torch.sum(1 / rank))\n", " total_count += count\n", " model_valid.detachFromDevice()\n", " mrr_per_epoch.append(mrr / total_count)\n", " print(f\"Epoch {epoch}, validation MRR {mrr / total_count:.4}\")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "bdf25049-234b-473c-92df-8492c8cba713", "metadata": {}, "source": [ "### Running inference on the trained model\n", "\n", "Finally, we can use our trained model to perform inference on FB15k-237." ] }, { "cell_type": "code", "execution_count": null, "id": "1ff1eae5-4f79-4061-b164-5c00be5bc01b", "metadata": {}, "outputs": [], "source": [ "inference_opts = poptorch.Options()\n", "inference_opts.setExecutionStrategy(poptorch.ShardedExecution(*pipeline_plan))\n", "inference_opts.enableExecutableCaching(executable_cache_dir)\n", "model_inference = poptorch.inferenceModel(model, options=inference_opts)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "2ad61176", "metadata": {}, "source": [ "We wrap the detaset with a `Prediction` object to simplify the inference process for all the different tasks." ] }, { "cell_type": "code", "execution_count": null, "id": "325d54ac-4cef-4e8d-a276-27b41b5cb54a", "metadata": {}, "outputs": [], "source": [ "pred = inference_utils.Prediction(\n", " dataset_inference[0],\n", " \"static/fb15k-237_entitymapping.txt\",\n", " osp.join(dataset_directory, \"FB15k-237/raw/\"),\n", ")" ] }, { "attachments": {}, "cell_type": "markdown", "id": "c7d3a99a", "metadata": {}, "source": [ "And run inference with this class." ] }, { "cell_type": "code", "execution_count": null, "id": "57d242f3", "metadata": {}, "outputs": [], "source": [ "pred.inference(model_inference, \"Good Will Hunting\", \"genre\", top_k=5)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "61445cf9-2575-409e-ad6c-8e6d4ba2aed5", "metadata": { "tags": [] }, "source": [ "## Running inference on the FB15k-237 graph\n", "\n", "Now it is time to test the model on the bigger FB15k-237 graph and make some predictions of the form `(head, relation, ?)`.\n", "We use a simple string comparison to match input strings to graph entities and relations. `pred.entity_vocab` and `pred.relation_vocab` contain lists of all available entities and relations.\n", "\n", "Note that the FB15k-237 graph is relatively small and not only lacks edges (which could be inferred using a knowledge graph completion model like this one) but also entities.\n", "\n", "`pred.inference` returns a list of entities and respective scores. Tails that occur in the graph are marked with an asterisk." ] }, { "cell_type": "code", "execution_count": null, "id": "e6599f01-fdeb-46db-9472-03af82f31194", "metadata": {}, "outputs": [], "source": [ "pred.inference(model_inference, \"London\", \"/location/location/contains\", top_k=5)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "195b7c89-632f-4cf4-b991-8bb826db22ab", "metadata": {}, "source": [ "## Interpreting the results\n", "\n", "Another advantage of the NBFNet model is its interpretability. By passing edge weights of `1.0` along all edges we can later compute the derivative of a prediction with respect to these weights and thus identify the paths that were most important for the prediction:" ] }, { "cell_type": "code", "execution_count": null, "id": "520e51d7-0d32-4e61-b938-29d57c622cdb", "metadata": {}, "outputs": [], "source": [ "pred.path_importance(model, head_id=4695, tail_id=5180, relation_id=31)" ] }, { "cell_type": "code", "execution_count": null, "id": "808b125e-591c-4e1b-ad3a-3f6768458b76", "metadata": {}, "outputs": [], "source": [ "pred.path_importance(model, head_id=12481, tail_id=1810, relation_id=4)" ] }, { "attachments": {}, "cell_type": "markdown", "id": "365cfd3a-fc93-42f1-8099-63ee810ee568", "metadata": {}, "source": [ "## Conclusion\n", "Using a subgraph of FB15k-237 we have trained an inductive link prediction model for knowledge graphs. This model has been used to infer missing connections in the full FB15k-237 graph and could demonstrate the applied reasoning by outputting the paths in the graph that were most relevant to a given prediction.\n", "\n", "As a next step you could try to speed up training by replicating the model four times on a POD-16 or train on a larger graph. this could be achieved by reducing the batch size or pipelining the model over more IPUs. \n", "\n", "If you are interested in node-level or graph-level tasks, take a look at our other examples. For instance, [Prediction of Molecular Properties using SchNet on Graphcore IPUs](../../schnet/pytorch_geometric/molecular_property_prediction_with_schnet.ipynb) for graph prediction or [Cluster GCN on IPU: Node classification task on a large graph using sampling](../../cluster_gcn/pytorch_geometric/node_classification_with_cluster_gcn.ipynb) for node prediction." ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }