{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Offline Tuple Packing for Graphs to reduce padding\n", "\n", "When processing a batch of graphs in machine learning models, it is common \n", "to combine (pack) several small graphs into one\n", "overall graph to accelerate processing and reduce the overhead of padding.\n", "This tutorial guides you through some data analysis on Graphs and graph packing to determine the efficient strategies for processing datasets drawn out of the OGB competition.\n", "\n", "It was shown in \"Packing: Towards 2x NLP BERT Acceleration. arXiv, jun 2021.\" by Matej Kosec et al. 2021 that for transformers padding can be a major part of the data and can cause a slow down of 50% and more. \n", "As a consequence, the distribution of sizes and degrees of graphs in datasets and how much padding they require, is an important consideration.\n", "\n", "This tutorial reproduces the results from the paper \"Tuple Packing: Efficient Batching of Small Graphs in Graph Neural Networks\" available on [arxiv](https://arxiv.org/abs/2209.06354).\n", "\n", "## Analysing graph sizes\n", "\n", "For the graph analysis, we focus on datasets in the OGB challenge that we load into PyTorch Geometric." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "import ogb\n", "from ogb.graphproppred import PygGraphPropPredDataset\n", "from torch_geometric.loader import DataLoader\n", "from ogb.utils.mol import smiles2graph\n", "\n", "ogb.utils.smiles2graph = smiles2graph\n", "from ogb.lsc.pcqm4mv2_pyg import PygPCQM4Mv2Dataset\n", "from pprint import pprint" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "\n", "This first step loads those datasets ready for additional processing:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "ogb_data = {}\n", "\n", "\n", "def get_training_dataloader(d_name: str) -> DataLoader:\n", " if d_name == \"ogbg-pcqm4mv2\":\n", " dataset = PygPCQM4Mv2Dataset(smiles2graph=smiles2graph)\n", " else:\n", " dataset = PygGraphPropPredDataset(name=d_name)\n", "\n", " split_idx = dataset.get_idx_split()\n", " train_loader = DataLoader(dataset[split_idx[\"train\"]], batch_size=1, shuffle=False)\n", " return train_loader\n", "\n", "\n", "for key in [\"ogbg-molhiv\", \"ogbg-molpcba\", \"ogbg-code2\", \"ogbg-pcqm4mv2\", \"ogbg-ppa\"]:\n", " print(\"loading:\", key)\n", " ogb_data[key] = get_training_dataloader(key)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We are interested in analysing the distribution in the number of edges and the number of nodes which is observed in the OGB small graph datasets.\n", "For this analysis, we do not need the graphs, we can limit ourselves to 2D histograms that map a number of graphs to tuples with the number of nodes and edges to get a sense of the size distribution of the data.\n", "\n", "Calculating the histograms can be a slow process so the results are cached to make running subsequent analysis faster." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from collections import defaultdict\n", "import pickle\n", "from tqdm import tqdm\n", "\n", "\n", "def get_histogram(data_loader):\n", " histogram = defaultdict(int)\n", " for item in tqdm(data_loader):\n", " histogram[(item.num_nodes, item.num_edges)] += 1\n", " return histogram\n", "\n", "\n", "load_histogram = True\n", "histogram_file = \"histograms.pkl\"\n", "\n", "if os.path.exists(histogram_file) and load_histogram:\n", " # Generating the histograms takes ~30mn so we load a saved version\n", " # of the histograms if it has already been generated\n", " with open(histogram_file, \"rb\") as f:\n", " ogb_histograms = pickle.load(f)\n", "else:\n", " ogb_histograms = {}\n", " for key in ogb_data:\n", " print(\"creating histogram:\", key)\n", " ogb_histograms[key] = get_histogram(ogb_data[key])\n", "\n", "# Save the histogram after processing\n", "with open(histogram_file, \"wb\") as f:\n", " pickle.dump(ogb_histograms, f)\n", "\n", "print(\"Displaying Histogram for 'ogbg-molhiv'\")\n", "print(ogb_histograms[\"ogbg-molhiv\"])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's first get some basic stats on the datasets:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_num_samples(histogram):\n", " return sum([histogram[key] for key in histogram])\n", "\n", "\n", "print(\"Number of graphs in each dataset:\")\n", "for key in ogb_histograms:\n", " print(key, get_num_samples(ogb_histograms[key]))" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_max_tuples_length(histogram):\n", " \"\"\"Get the maximum entry size for each tuple component\"\"\"\n", " maximum_length = []\n", " for key in histogram:\n", " if not maximum_length:\n", " maximum_length = list(key)\n", " for index, entry in enumerate(maximum_length):\n", " maximum_length[index] = max(entry, key[index])\n", " return maximum_length\n", "\n", "\n", "# getting max_tuples_length\n", "ogbg_mtl_dict = {}\n", "for key in ogb_histograms:\n", " ogbg_mtl_dict[key] = get_max_tuples_length(ogb_histograms[key])\n", "print(\"Maximum number of nodes and edges in each graph:\")\n", "pprint(ogbg_mtl_dict)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Detailed statistics on padding efficiency will be provided later in the tutorial.\n", "\n", "For now, let's explore the distribution of number of nodes and edges in each of the datasets:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import matplotlib\n", "import matplotlib.pyplot as plt\n", "import os\n", "from pathlib import Path\n", "\n", "figure_directory = Path(\"graph_packing_tutorial\")\n", "figure_directory.mkdir(parents=True, exist_ok=True)\n", "\n", "\n", "def visualize_2D_histogram(histogram, key, dropout=0.01):\n", " total_count = sum([histogram[(nodes, edges)] for nodes, edges in histogram])\n", " threshold = total_count * dropout / 100\n", " num_nodes = [\n", " nodes for nodes, edges in histogram if histogram[(nodes, edges)] >= threshold\n", " ]\n", " num_edges = [\n", " edges for nodes, edges in histogram if histogram[(nodes, edges)] >= threshold\n", " ]\n", " image = np.zeros([max(num_nodes) + 1, max(num_edges) + 1])\n", " exceptions = []\n", " for nodes, edges in histogram:\n", " try:\n", " image[nodes][edges] = histogram[(nodes, edges)]\n", " except IndexError:\n", " exceptions.append((nodes, edges, histogram[(nodes, edges)]))\n", " if exceptions:\n", " print(\n", " \"Not visualised:\", sum([i[2] for i in exceptions]) / total_count * 100, \"%\"\n", " )\n", " fig = plt.figure(dpi=200)\n", " im = plt.imshow(image, cmap=plt.cm.hot_r)\n", " cb = plt.colorbar(shrink=0.5)\n", " cb.set_label(\"Number of samples\")\n", " plt.xlabel(\"Number of Edges\")\n", " plt.ylabel(\"Number of Nodes\")\n", " plt.title(\"Dataset: \" + key)\n", " fig.savefig(figure_directory / f\"{key}_dual_histogram.png\", bbox_inches=\"tight\")\n", "\n", "\n", "for key in ogb_histograms:\n", " if not key == \"ogbg-ppa\":\n", " print(\"visualizing 2D histogram:\", key)\n", " visualize_2D_histogram(ogb_histograms[key], key, dropout=0.01)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def visualize_2D_histogram_builtin(histogram, key, dropout=0.01):\n", " \"\"\"Uses plt.hist2d to do the same as the figure above\"\"\"\n", " total_count = sum([histogram[(nodes, edges)] for nodes, edges in histogram])\n", " threshold = total_count * dropout / 100\n", " raw_num_nodes = [nodes for nodes, edges in histogram]\n", " raw_num_edges = [edges for nodes, edges in histogram]\n", " raw_num_weights = [weights for weights in histogram.values()]\n", " fig = plt.figure(dpi=200)\n", " im = plt.hist2d(\n", " raw_num_edges,\n", " raw_num_nodes,\n", " weights=raw_num_weights,\n", " cmap=plt.cm.hot_r,\n", " bins=max(raw_num_nodes),\n", " cmin=threshold,\n", " )\n", " plt.xlabel(\"Number of Edges\")\n", " plt.ylabel(\"Number of Nodes\")\n", " plt.title(\"Dataset: \" + key)\n", " cb = plt.colorbar(shrink=0.5)\n", " cb.set_label(\"Number of samples\")\n", " fig.savefig(figure_directory / f\"{key}_dual_histogram_v2.png\", bbox_inches=\"tight\")\n", "\n", "\n", "for key in ogb_histograms:\n", " if not key == \"ogbg-ppa\":\n", " print(\"visualizing 2D histogram:\", key)\n", " visualize_2D_histogram_builtin(ogb_histograms[key], key, dropout=0.01)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In the histograms, we focus on major bins and ignore the tails to get a better and faster visualisation.\n", "For the `ogbg-ppa` this results in an empty image.\n", "Feel free to explore different drop out rates and the `ogbg-ppa` dataset.\n", "You can also explore the histograms for the number of nodes or edges." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def visualize_nodes_edges_histogram(histogram, key, cumulative=False):\n", " num_nodes = []\n", " num_edges = []\n", " for nodes, edges in histogram:\n", " count = histogram[(nodes, edges)]\n", " num_nodes.extend([nodes] * count)\n", " num_edges.extend([edges] * count)\n", "\n", " min_num_nodes = int(min(num_nodes))\n", " max_num_nodes = int(max(num_edges))\n", "\n", " plt.hist(num_nodes, max_num_nodes - min_num_nodes + 1, cumulative=cumulative)\n", " plt.xlabel(\"Number of Nodes\")\n", " plt.ylabel(\"Counts\")\n", " plt.title(\"Histogram of number of graph nodes for \" + key)\n", " plt.savefig(\n", " \"graph_packing_tutorial\"\n", " + os.sep\n", " + (\"cumulative\" if cumulative else \"\")\n", " + key\n", " + \"_nodes_histogram.png\",\n", " bbox_inches=\"tight\",\n", " )\n", " plt.show()\n", "\n", " min_num_edges = int(min(num_edges))\n", " max_num_edges = int(max(num_edges))\n", "\n", " plt.hist(\n", " num_edges, max_num_edges - min_num_edges + 1, cumulative=cumulative\n", " ) # max_num_edges-min_num_edges+1)\n", " plt.xlabel(\"Number of Edges\")\n", " plt.ylabel(\"Counts\")\n", " plt.title(f\"Histogram of number of edges for \" + key)\n", " plt.savefig(\n", " \"graph_packing_tutorial\"\n", " + os.sep\n", " + (\"cumulative\" if cumulative else \"\")\n", " + key\n", " + \"_edges_histogram.png\",\n", " bbox_inches=\"tight\",\n", " )\n", " plt.show()\n", "\n", "\n", "for key in ogb_histograms:\n", " print(\"visualizing histograms: \", key)\n", " visualize_nodes_edges_histogram(ogb_histograms[key], key, cumulative=False)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "From the visualizations, it is evident that most graphs are much smaller than the largest graph in the dataset: this means that, for a batch size of 1, a major portion of the data will be padding.\n", "\n", "For larger batch sizes, data usually gets combined which on average may lead to a better distribution, however, the size of a batch in theory can vary as much as we have variance between individual graphs. To minimise padding, it makes sense to batch graphs that result in a constant size of nodes and edges when combined." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Tuple Packing\n", "\n", "To minimise the amount of padding needed to process a dataset we develop a heuristic which assigns a \"sorting priority\" to each graph and we will then assemble *packs* based on this priority.\n", "This heuristic decides which graphs to tackle first as well as which graphs to prefer when filling empty spots:\n", "a graph with a higher priority should be assigned to a *pack* before a graph with a lower priority. \n", "There is a variety of possibilities for this heuristic. We provide the most intuitive ones." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "heuristics = {\n", " \"prod\": lambda x, y: int(x * y),\n", " \"sum\": lambda x, y: int(x + y),\n", " \"max\": lambda x, y: max(x, y),\n", " \"min\": lambda x, y: max(x, y),\n", " \"node\": lambda x, y: int(x),\n", " \"edge\": lambda x, y: int(y),\n", "}" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now, we need some fast algorithm that decides which graphs to combine based on the heuristic.\n", "\n", "Remember that the number of graphs can be large, hence, we have to operate on the histogram to get a solution fast.\n", "\n", "Our approach is at its core to sort the histogram by the heuristic and iterate over it.\n", "At the beginning as well if we can't find any pack, where we can add the graph and still meet the size constraints, we start a new pack with the current graph.\n", "A pack describes which tuple sizes get combined and how many of these combinations are obtained.\n", "For each pack, we calculate how much space is left in each component and apply the heuristic for sorting. \n", "If multiple packs would fit a new incoming graph, we add it to the one with the smallest heuristic. This way, we get a near optimal fit.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def pack_using_dlpfhp(\n", " histogram,\n", " max_tuple_length,\n", " max_tuples_per_pack,\n", " heuristic=lambda x, y: int(x * y),\n", " verbose=True,\n", "):\n", " \"\"\"Dual Longest-pack-first histogram-packing algorithm.\n", "\n", " Arguments:\n", " histogram Dict[Tuple[int, int], int]: The histogram of the dataset, it maps\n", " pairs of node and edge numbers to the number of graphs which match this\n", " specific size in the dataset.\n", " max_tuple_length (Tuple[int, int]): A pair that describes the maximum size of\n", " the container for each component that must be filled with\n", " packing. In this example this is a maximum number of nodes or edges.\n", " max_tuples_per_pack (int | Literal[\"max\"]): This integer parameter limits how\n", " many tuples/graphs can be combined. If using \"max\", no limit on packs is\n", " set, which in some cases can slow down the packing algorithm drastically.\n", " heuristic (Callable[int, int]): A function which calculates the priority heuristic\n", " from the histogram key.\n", " \"\"\"\n", " # heuristic assignment\n", " heuristic_data_list = [\n", " (heuristic(a, b), a, b, histogram[(a, b)]) for a, b in histogram\n", " ]\n", " heuristic_data_list.sort()\n", " heuristic_data_list.reverse()\n", " data_list = heuristic_data_list\n", " max_a, max_b = max_tuple_length[0], max_tuple_length[1]\n", " max_size = heuristic(max_a, max_b)\n", " if max_tuples_per_pack == \"max\":\n", " max_tuples_per_pack = min(max_tuple_length)\n", " # Initialize main strategy data dictionary.\n", " # The key indicates how much space is left.\n", " # The value is a list of tuples, consisting of counts and respective packs/tuples.\n", " tmp_strategies_per_length = defaultdict(list)\n", " strategies_per_length = defaultdict(list)\n", " for i in range(len(data_list)): # distribute each bin of histogram\n", " size, len_a, len_b, n_sequences_to_bin = data_list[i]\n", " left_size = heuristic(max_a - len_a, max_b - len_b)\n", " offset = 0 # smallest possible offset for perfect fit\n", " while n_sequences_to_bin > 0:\n", " keys = [key for key in tmp_strategies_per_length if key >= size + offset]\n", " if not keys:\n", " offset = max_size + 1\n", " else:\n", " offset = min(keys) - size\n", " if (size + offset) in tmp_strategies_per_length:\n", " for i in range(len(tmp_strategies_per_length[size + offset])):\n", " lens_a, lens_b, n_sequences_to_pack = tmp_strategies_per_length[\n", " size + offset\n", " ][i]\n", " if (len_a + sum(lens_a)) <= max_a and (\n", " len_b + sum(lens_b)\n", " ) <= max_b:\n", " tmp_strategies_per_length[size + offset].pop(i)\n", " new_lens_a = lens_a.copy()\n", " new_lens_a.append(len_a)\n", " new_lens_b = lens_b.copy()\n", " new_lens_b.append(len_b)\n", " new_size = heuristic(\n", " max_a - sum(new_lens_a), max_b - sum(new_lens_b)\n", " )\n", " new_count = min(n_sequences_to_pack, n_sequences_to_bin)\n", " # adjust strategies\n", " if n_sequences_to_pack > new_count:\n", " tmp_strategies_per_length[size + offset].append(\n", " (lens_a, lens_b, n_sequences_to_pack - new_count)\n", " )\n", " if not tmp_strategies_per_length[size + offset]:\n", " tmp_strategies_per_length.pop(size + offset)\n", " if new_size == 0 or max_tuples_per_pack == len(new_lens_a):\n", " strategies_per_length[0].append(\n", " (new_lens_a, new_lens_b, new_count)\n", " )\n", " else:\n", " tmp_strategies_per_length[new_size].append(\n", " (new_lens_a, new_lens_b, new_count)\n", " )\n", " n_sequences_to_bin -= new_count\n", " offset = 0\n", " break\n", " offset += 1\n", " if offset + size > max_size:\n", " new_size = heuristic(max_a - len_a, max_b - len_b)\n", " if new_size == 0 or max_tuples_per_pack == 1:\n", " strategies_per_length[new_size].append(\n", " ([len_a], [len_b], n_sequences_to_bin)\n", " )\n", " else:\n", " tmp_strategies_per_length[new_size].append(\n", " ([len_a], [len_b], n_sequences_to_bin)\n", " )\n", " n_sequences_to_bin = 0\n", " break\n", " # merge all strategies\n", " for key in tmp_strategies_per_length:\n", " strategies_per_length[key].extend(tmp_strategies_per_length[key])\n", " # flatten strategies dictionary\n", " strategy_set = []\n", " strategy_repeat_count = []\n", " sum_lens_a, sum_lens_b = [], []\n", " for key in strategies_per_length:\n", " for lens_a, lens_b, n_sequences_to_pack in strategies_per_length[key]:\n", " strategy_set.append((lens_a, lens_b))\n", " strategy_repeat_count.append(n_sequences_to_pack)\n", " sum_lens_a.append(sum(lens_a))\n", " sum_lens_b.append(sum(lens_b))\n", " if not (max_a == max(sum_lens_a) and max_b == max(sum_lens_b)):\n", " if verbose:\n", " print(\n", " \"max discrepancy, reducing sequence length\",\n", " max_a,\n", " max(sum_lens_a),\n", " max_b,\n", " max(sum_lens_b),\n", " )\n", " max_a, max_b = max(sum_lens_a), max(sum_lens_b)\n", " # efficiency calculation\n", " empty_tokens_a = int(\n", " sum(\n", " [\n", " count * (max_a - sum(pack_a))\n", " for count, (pack_a, pack_b) in zip(strategy_repeat_count, strategy_set)\n", " ]\n", " )\n", " )\n", " empty_tokens_b = int(\n", " sum(\n", " [\n", " count * (max_b - sum(pack_b))\n", " for count, (pack_a, pack_b) in zip(strategy_repeat_count, strategy_set)\n", " ]\n", " )\n", " )\n", " packs = int(sum(strategy_repeat_count))\n", " total_tokens_a, total_tokens_b = int(max_a * packs), int(max_b * packs)\n", " token_efficiency = (\n", " 100 - empty_tokens_a / total_tokens_a * 100,\n", " 100 - empty_tokens_b / total_tokens_b * 100,\n", " )\n", " return strategy_set, np.array(strategy_repeat_count), token_efficiency" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Note that the code is written for tuples with two components like the number of nodes and edges but can easily be extended to more components.\n", "\n", "The packing function calculates the \"token efficiency\" of the heuristic: it is the ratio of true data over the size of the padded data; larger numbers are better.\n", "\n", "Now let's try a few heuristics to get a better picture." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "max_items_per_pack = 256\n", "# other possibilities: \"max\", 1\n", "heuristic = lambda x, y: int(x * y)\n", "\n", "for key in ogb_histograms:\n", " # \"ogbg-ppa\" processing takes too long\n", " if key == \"ogbg-ppa\":\n", " continue\n", " print(\"----------------------------------------------------\")\n", " print(f\"Token efficiencies for '{key}' with:\")\n", " print(f\" - maximum container sizes={ogbg_mtl_dict[key]}\")\n", " print(f\" - max number of graphs per pack={max_items_per_pack}\")\n", " print(\" strategy, nodes, edges\\n ------------------------\")\n", " for heuristic in heuristics:\n", " _, _, token_efficiency = pack_using_dlpfhp(\n", " ogb_histograms[key],\n", " ogbg_mtl_dict[key],\n", " max_items_per_pack,\n", " heuristic=heuristics[heuristic],\n", " )\n", " print(f\" {heuristic}, {token_efficiency[0]:.2f}%, {token_efficiency[1]:.2f}%\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "After comparing `max_items_per_pack=1` and `max_items_per_pack=256`, you will see that the efficiency improves drastically. You can also use the first case to get an estimate of the speedup you can achieve by combining graphs." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Tuning the pack sizes\n", "\n", "The number of tokens contained in each pack can be changed with the `max_tuple_length` argument.\n", "It is an important parameter which can be tuned to improve the token efficiency of the packed dataset.\n", "While the minimum size of the packs must be at least the size of the largest graph in the unpacked dataset, it can be made larger to allow more graphs to be assembled together for a higher token efficiency.\n", "\n", "We will explore the effect of the `max_tuple_length` argument on the token efficiency when packing the `\"ogbg-pcqm4mv2\"` dataset.\n", "\n", "For each parameter we collect the \"node efficiency\" and the \"edge efficiency\" which is the ratio of useful tokens over the total space for tokens in the packed dataset. These metrics give a measure of how much compute will be wasted on padding data.\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from tqdm import tqdm\n", "\n", "dataset_to_analyse = \"ogbg-pcqm4mv2\"\n", "\n", "max_nodes, max_edges = ogbg_mtl_dict[dataset_to_analyse]\n", "histogram = ogb_histograms[dataset_to_analyse]\n", "results = []\n", "for pack_size_a in tqdm(range(max_nodes, 3 * max_nodes, 1)):\n", " for pack_size_b in range(max_edges, 3 * max_edges, 2):\n", " strategy_set, strategy_repeat_count, eff = pack_using_dlpfhp(\n", " histogram,\n", " max_tuple_length=[pack_size_a, pack_size_b],\n", " max_tuples_per_pack=\"max\",\n", " heuristic=lambda x, y: x * y,\n", " verbose=False,\n", " )\n", " results.append(((pack_size_a, pack_size_b), eff))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To compare our parameters we calculate the harmonic mean of the node and edge efficiencies, giving us a single metric representing the performance of a parameter combination.\n", "Let's sort our results to find the best parameter combinations." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import pandas as pd\n", "\n", "harmonic_efficiency = [\n", " (2 / (1 / node_eff + 1 / edge_eff), node_eff, edge_eff, *sizes)\n", " for sizes, (node_eff, edge_eff) in results\n", "]\n", "harmonic_efficiency.sort()\n", "harmonic_efficiency.reverse()\n", "efficiency_df = pd.DataFrame(\n", " harmonic_efficiency,\n", " columns=[\n", " \"harmonic_eff\",\n", " \"node_eff\",\n", " \"edge_eff\",\n", " \"node_pack_size\",\n", " \"edge_pack_size\",\n", " ],\n", ")\n", "efficiency_df.set_index(keys=[\"node_pack_size\", \"edge_pack_size\"], inplace=True)\n", "efficiency_df.head()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now let us compare to the results for the minimal value of pack size:" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "efficiency_df.loc[[ogbg_mtl_dict[dataset_to_analyse]], :]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "To gain a better understanding of the distribution of where efficient packing parameters cluster, let us visualize the harmonic mean of the efficiencies for all combinations of parameters." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import numpy\n", "\n", "result_image = numpy.zeros([max_edges, 2 * max_nodes])\n", "edges = []\n", "nodes = []\n", "\n", "for (n, e), d in results:\n", " result_image[(e - max_edges) // 2][n - max_nodes] = max(\n", " 50, 2 / (1 / d[0] + 1 / d[1])\n", " )\n", " edges.append(e)\n", " nodes.append(n)\n", "\n", "\n", "fig = plt.figure(dpi=300)\n", "extent = None\n", "extent = np.min(edges), np.max(edges), np.max(nodes), np.min(nodes)\n", "im = plt.imshow(result_image.T, cmap=plt.cm.hot_r, extent=extent)\n", "cb = plt.colorbar(shrink=0.5)\n", "cb.set_label(\"Harmonic Efficiency (%)\")\n", "plt.xlabel(\"Number of Edges\")\n", "plt.ylabel(\"Number of Nodes\")\n", "fig.savefig(\n", " \"graph_packing_tutorial\" + os.sep + \"PCQM4Mv2_tuple_pack_performance.png\",\n", " bbox_inches=\"tight\",\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can see that the best harmonic mean of the packing efficiency on the nodes and edges is obtained at 30 nodes and 62 edges with a value of (98.6%, 99.0%).\n", "It is close to perfect, and much better than the 70% achieved when using the maximum sizes found in the ogbg-pcqm4mv2 dataset (20 nodes and 56 edges).\n", "We also get good results for 45 nodes and 96 edges (94.9%, 98.4%). Having multiple efficient packing parameter combinations is useful as it lets us tradeoff some packing efficiency with compute efficiency, as bigger packs may allow us to use the IPUs more efficiently, while smaller packs can help a model fit on fewer IPUs. The good results are mostly centred close to the diagonal where there are twice as many edges as nodes." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Conclusion\n", "\n", "This tutorial is a companion to the paper \"Tuple Packing: Efficient Batching of Small Graphs in Graph Neural Networks\" available on [arxiv](https://arxiv.org/abs/2209.06354).\n", "It worked through the calculation of efficient packs for a range of datasets in the OGB competition.\n", "We showed how histograms can be used to efficiently calculate heuristics and assemble optimised packs of graphs for processing in GNNs.\n", "We also analysed the importance of pack size on the efficiency of the packed dataset.\n", "\n", "If you like this tutorial and use the code, please cite our paper." ] }, { "cell_type": "markdown", "metadata": {}, "source": [] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 4 }