GNN/7_Aggregation_Package.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Install Pytorch and PyG"
      ],
      "metadata": {
        "id": "6tFC-bQ87Amb"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zWn5yzT0LOzH"
      },
      "source": [
        "# Install required packages.\n",
        "import os\n",
        "import torch\n",
        "os.environ['TORCH'] = torch.__version__\n",
        "\n",
        "!pip install -q torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html\n",
        "!pip install -q torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html\n",
        "!pip install -q git+https://github.com/pyg-team/pytorch_geometric.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KDy46FIQ6OWN"
      },
      "source": [
        "# Customizing Aggregations within Message Passing\n",
        "\n",
        "Aggregation functions play an important role in the message passing framework and the readout function when implementing GNNs. Many works in the GNN literature ([Hamilton et al. (2017)](https://cs.stanford.edu/~jure/pubs/graphsage-nips17.pdf), [Xu et al. (2018)](https://arxiv.org/abs/1810.00826), [Corso et al. (2020)](https://proceedings.neurips.cc/paper/2020/file/99cad265a1768cc2dd013f0e740300ae-Paper.pdf), [Li et al. (2020)](https://arxiv.org/abs/2006.07739)), demonstrate that the choice of aggregation functions contributes significantly to the performance of GNN models. In particular, the performance of GNNs with different aggregation functions differs when applied to distinct tasks and datasets. Recent works also show that using multiple aggregations ([Corso et al. (2020)](https://proceedings.neurips.cc/paper/2020/file/99cad265a1768cc2dd013f0e740300ae-Paper.pdf)) and learnable aggregations ([Li et al. (2020)](https://arxiv.org/abs/2006.07739)) can potentially gain substantial improvements. To facilitate experimentation with these different aggregation schemes and unify concepts of aggregation within GNNs across both [`MessagePassing`](https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/conv/message_passing.py) and [global readouts](https://github.com/pyg-team/pytorch_geometric/tree/master/torch_geometric/nn/glob), we provide **modular and re-usable aggregations** in the newly defined `torch_geometric.nn.aggr.*` package. Unifying these concepts also helps us to perform optimization and specialized implementations in a single place. In the new integration, the following functionality is applicable:\n",
        "\n",
        "```python\n",
        "# Original interface with string type as aggregation argument\n",
        "class MyConv(MessagePassing):\n",
        "    def __init__(self):\n",
        "        super().__init__(aggr=\"mean\")\n",
        "\n",
        "\n",
        "# Use a single aggregation module as aggregation argument\n",
        "class MyConv(MessagePassing):\n",
        "    def __init__(self):\n",
        "        super().__init__(aggr=MeanAggregation())\n",
        "\n",
        "\n",
        "# Use a list of aggregation strings as aggregation argument\n",
        "class MyConv(MessagePassing):\n",
        "    def __init__(self):\n",
        "        super().__init__(aggr=['mean', 'max', 'sum', 'std', 'var'])\n",
        "\n",
        "\n",
        "# Use a list of aggregation modules as aggregation argument\n",
        "class MyConv(MessagePassing):\n",
        "    def __init__(self):\n",
        "        super().__init__(aggr=[\n",
        "            MeanAggregation(),\n",
        "            MaxAggregation(),\n",
        "            SumAggregation(),\n",
        "            StdAggregation(),\n",
        "            VarAggregation(),\n",
        "        ])\n",
        "\n",
        "\n",
        "# Use a list of mixed modules and strings as aggregation argument\n",
        "class MyConv(MessagePassing):\n",
        "    def __init__(self):\n",
        "        super().__init__(aggr=[\n",
        "            'mean',\n",
        "            MaxAggregation(),\n",
        "            'sum',\n",
        "            StdAggregation(),\n",
        "            'var',\n",
        "        ])\n",
        "\n",
        "\n",
        "# Define multiple aggregations with `MultiAggregation` module\n",
        "class MyConv(MessagePassing):\n",
        "    def __init__(self):\n",
        "        super().__init__(aggr=MultiAggregation([\n",
        "            SoftmaxAggregation(t=0.1, learn=True),\n",
        "            SoftmaxAggregation(t=1, learn=True),\n",
        "            SoftmaxAggregation(t=10, learn=True)\n",
        "        ]))\n",
        "\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this tutorial, we explore the new aggregation package with `SAGEConv` ([Hamilton et al. (2017)](https://cs.stanford.edu/~jure/pubs/graphsage-nips17.pdf)) and `ClusterLoader` ([Chiang et al. (2019)](https://arxiv.org/abs/1905.07953)) and showcase on the `PubMed` graph from the `Planetoid` node classification benchmark suite ([Yang et al. (2016)](https://arxiv.org/abs/1603.08861))."
      ],
      "metadata": {
        "id": "3W4IfAWgSXVy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Loading the dataset\n",
        "Let's first load the `Planetoid` dataset and create subgraphs with `ClusterData` for training."
      ],
      "metadata": {
        "id": "r_s25N9oyBIy"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "eBN2pGDueDpZ",
        "outputId": "7139ed04-cf1b-4406-80e3-f6b3ac10f05c"
      },
      "source": [
        "import torch\n",
        "from torch_geometric.datasets import Planetoid\n",
        "from torch_geometric.transforms import NormalizeFeatures\n",
        "\n",
        "dataset = Planetoid(root='data/Planetoid', name='PubMed',\n",
        "                    transform=NormalizeFeatures())\n",
        "\n",
        "print()\n",
        "print(f'Dataset: {dataset}:')\n",
        "print('==================')\n",
        "print(f'Number of graphs: {len(dataset)}')\n",
        "print(f'Number of features: {dataset.num_features}')\n",
        "print(f'Number of classes: {dataset.num_classes}')\n",
        "\n",
        "data = dataset[0]  # Get the first graph object.\n",
        "\n",
        "print()\n",
        "print(data)\n",
        "\n",
        "from torch_geometric.loader import ClusterData, ClusterLoader\n",
        "\n",
        "seed = 42\n",
        "torch.manual_seed(seed)\n",
        "cluster_data = ClusterData(data, num_parts=128)  # 1. Create subgraphs.\n",
        "train_loader = ClusterLoader(cluster_data, batch_size=32,\n",
        "                             shuffle=True)  # 2. Stochastic partioning scheme."
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Dataset: PubMed():\n",
            "==================\n",
            "Number of graphs: 1\n",
            "Number of features: 500\n",
            "Number of classes: 3\n",
            "\n",
            "Data(x=[19717, 500], edge_index=[2, 88648], y=[19717], train_mask=[19717], val_mask=[19717], test_mask=[19717])\n",
            "===============================================================================================================\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Computing METIS partitioning...\n",
            "Done!\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Define train, test, and run functions\n",
        "Here we define a simple `run` function for training the GNN model."
      ],
      "metadata": {
        "id": "tbHL1x7fzjeR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "criterion = torch.nn.CrossEntropyLoss()\n",
        "\n",
        "\n",
        "def train(model):\n",
        "    model.train()\n",
        "    optimizer = torch.optim.Adam(model.parameters(), lr=0.01,\n",
        "                                 weight_decay=5e-4)\n",
        "    for sub_data in train_loader:  # Iterate over each mini-batch.\n",
        "        optimizer.zero_grad()  # Clear gradients.\n",
        "        out = model(sub_data.x,\n",
        "                    sub_data.edge_index)  # Perform a single forward pass.\n",
        "        loss = criterion(\n",
        "            out[sub_data.train_mask], sub_data.y[sub_data.train_mask]\n",
        "        )  # Compute the loss solely based on the training nodes.\n",
        "        loss.backward()  # Derive gradients.\n",
        "        optimizer.step()  # Update parameters based on gradients.\n",
        "\n",
        "\n",
        "def test(model):\n",
        "    model.eval()\n",
        "    out = model(data.x, data.edge_index)\n",
        "    pred = out.argmax(dim=1)  # Use the class with highest probability.\n",
        "\n",
        "    accs = []\n",
        "    for mask in [data.train_mask, data.val_mask, data.test_mask]:\n",
        "        correct = pred[mask] == data.y[\n",
        "            mask]  # Check against ground-truth labels.\n",
        "        accs.append(int(correct.sum()) /\n",
        "                    int(mask.sum()))  # Derive ratio of correct predictions.\n",
        "    return accs\n",
        "\n",
        "\n",
        "def run(model, epochs=5):\n",
        "    for epoch in range(epochs):\n",
        "        loss = train(model)\n",
        "        train_acc, val_acc, test_acc = test(model)\n",
        "        print(\n",
        "            f'Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val Acc: {val_acc:.4f}, Test Acc: {test_acc:.4f}'\n",
        "        )"
      ],
      "metadata": {
        "id": "5pJQ7brC7VzC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Define a GNN class and Import Aggregations\n",
        "Now, let's define a GNN helper class and import all those new aggregation operators!\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "ALbxt90iftTm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import copy\n",
        "\n",
        "import torch.nn.functional as F\n",
        "from torch_geometric.nn import (\n",
        "    Aggregation,\n",
        "    MaxAggregation,\n",
        "    MeanAggregation,\n",
        "    MultiAggregation,\n",
        "    SAGEConv,\n",
        "    SoftmaxAggregation,\n",
        "    StdAggregation,\n",
        "    SumAggregation,\n",
        "    VarAggregation,\n",
        ")\n",
        "\n",
        "\n",
        "class GNN(torch.nn.Module):\n",
        "    def __init__(self, hidden_channels, aggr='mean', aggr_kwargs=None):\n",
        "        super().__init__()\n",
        "        self.conv1 = SAGEConv(\n",
        "            dataset.num_node_features,\n",
        "            hidden_channels,\n",
        "            aggr=aggr,\n",
        "            aggr_kwargs=aggr_kwargs,\n",
        "        )\n",
        "        self.conv2 = SAGEConv(\n",
        "            hidden_channels,\n",
        "            dataset.num_classes,\n",
        "            aggr=copy.deepcopy(aggr),\n",
        "            aggr_kwargs=aggr_kwargs,\n",
        "        )\n",
        "\n",
        "    def forward(self, x, edge_index):\n",
        "        x = self.conv1(x, edge_index)\n",
        "        x = x.relu()\n",
        "        x = F.dropout(x, p=0.5, training=self.training)\n",
        "        x = self.conv2(x, edge_index)\n",
        "        return x"
      ],
      "metadata": {
        "id": "dhGCq1KbAxAX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Original interface with string type as the aggregation argument\n",
        "Previously, PyG only supports customizing [MessagePassing](https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/conv/message_passing.py) with simple aggregations (e.g., `'mean'`, `'max'`, `'sum'`). Let's define a GNN with `mean` aggregation and run it for 5 epochs."
      ],
      "metadata": {
        "id": "ZIg70vDmLO5V"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(seed)\n",
        "model = GNN(16, aggr='mean')\n",
        "print(model)\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)\n",
        "run(model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GuRdM5VjDIAv",
        "outputId": "c9ca7892-8f37-425f-e884-791a86916fef"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GNN(\n",
            "  (conv1): SAGEConv(500, 16, aggr=mean)\n",
            "  (conv2): SAGEConv(16, 3, aggr=mean)\n",
            ")\n",
            "Epoch: 000, Train: 0.6000, Val Acc: 0.4400, Test Acc: 0.4240\n",
            "Epoch: 001, Train: 0.6167, Val Acc: 0.5380, Test Acc: 0.5120\n",
            "Epoch: 002, Train: 0.6500, Val Acc: 0.5440, Test Acc: 0.5180\n",
            "Epoch: 003, Train: 0.7167, Val Acc: 0.5540, Test Acc: 0.5250\n",
            "Epoch: 004, Train: 0.9000, Val Acc: 0.7200, Test Acc: 0.6920\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Use a single aggregation module as the aggregation argument\n",
        "In the new interface, the [MessagePassing](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.MessagePassing) class can take an [Aggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.Aggregation) module as an argument. Here we can define the mean aggregation by `MeanAggregation`. We can see the model achieves the same performance as previously."
      ],
      "metadata": {
        "id": "C-e0uXR9Ldcf"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(seed)\n",
        "model = GNN(16, aggr=MeanAggregation())\n",
        "print(model)\n",
        "run(model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "C66kK8wfCSnF",
        "outputId": "b59369ad-c612-4ec3-c199-732339a656e8"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GNN(\n",
            "  (conv1): SAGEConv(500, 16, aggr=MeanAggregation())\n",
            "  (conv2): SAGEConv(16, 3, aggr=MeanAggregation())\n",
            ")\n",
            "Epoch: 000, Train: 0.6000, Val Acc: 0.4400, Test Acc: 0.4240\n",
            "Epoch: 001, Train: 0.6167, Val Acc: 0.5380, Test Acc: 0.5120\n",
            "Epoch: 002, Train: 0.6500, Val Acc: 0.5440, Test Acc: 0.5180\n",
            "Epoch: 003, Train: 0.7167, Val Acc: 0.5540, Test Acc: 0.5250\n",
            "Epoch: 004, Train: 0.9000, Val Acc: 0.7200, Test Acc: 0.6920\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Use a list of aggregation strings as the aggregation argument\n",
        "\n",
        "For defining multiple aggregations, we can use a list of strings as the input argument. The aggregations will be **resolved from pure strings** via a lookup table, following the design principles of the [class-resolver](https://github.com/cthoyt/class-resolver) library, e.g., by simply passing in `\"mean\"` to the [**MessagePassing**](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.MessagePassing) module. This will automatically resolve it to the MeanAggregation class. Let's see how a PNA-like GNN ([Corso et al. (2020)](https://proceedings.neurips.cc/paper/2020/file/99cad265a1768cc2dd013f0e740300ae-Paper.pdf)) works. It converges much faster!"
      ],
      "metadata": {
        "id": "yslL-7_vLmaD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(seed)\n",
        "model = GNN(16, aggr=['mean', 'max', 'sum', 'std', 'var'])\n",
        "print(model)\n",
        "run(model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5ANwoZ7Z-aau",
        "outputId": "4f2abae7-a71c-413d-9d7f-75caaa58c4d2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GNN(\n",
            "  (conv1): SAGEConv(500, 16, aggr=['mean', 'max', 'sum', 'std', 'var'])\n",
            "  (conv2): SAGEConv(16, 3, aggr=['mean', 'max', 'sum', 'std', 'var'])\n",
            ")\n",
            "Epoch: 000, Train: 0.7833, Val Acc: 0.6420, Test Acc: 0.6320\n",
            "Epoch: 001, Train: 0.8333, Val Acc: 0.7040, Test Acc: 0.6980\n",
            "Epoch: 002, Train: 0.8500, Val Acc: 0.6660, Test Acc: 0.6570\n",
            "Epoch: 003, Train: 0.9500, Val Acc: 0.7120, Test Acc: 0.7010\n",
            "Epoch: 004, Train: 0.9333, Val Acc: 0.7500, Test Acc: 0.7420\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Use a list of aggregation modules as the aggregation argument\n",
        "You can also use a list of [Aggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.Aggregation) modules to specify your convolutions."
      ],
      "metadata": {
        "id": "cykfCUUhLojw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(seed)\n",
        "model = GNN(\n",
        "    16, aggr=[\n",
        "        MeanAggregation(),\n",
        "        MaxAggregation(),\n",
        "        SumAggregation(),\n",
        "        StdAggregation(),\n",
        "        VarAggregation(),\n",
        "    ])\n",
        "print(model)\n",
        "run(model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0M1UdXiOGpuC",
        "outputId": "ad4aad9b-06e8-40fe-ebe7-6578655df300"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GNN(\n",
            "  (conv1): SAGEConv(500, 16, aggr=['MeanAggregation()', 'MaxAggregation()', 'SumAggregation()', 'StdAggregation()', 'VarAggregation()'])\n",
            "  (conv2): SAGEConv(16, 3, aggr=['MeanAggregation()', 'MaxAggregation()', 'SumAggregation()', 'StdAggregation()', 'VarAggregation()'])\n",
            ")\n",
            "Epoch: 000, Train: 0.7833, Val Acc: 0.6420, Test Acc: 0.6320\n",
            "Epoch: 001, Train: 0.8333, Val Acc: 0.7040, Test Acc: 0.6980\n",
            "Epoch: 002, Train: 0.8500, Val Acc: 0.6660, Test Acc: 0.6570\n",
            "Epoch: 003, Train: 0.9500, Val Acc: 0.7120, Test Acc: 0.7010\n",
            "Epoch: 004, Train: 0.9333, Val Acc: 0.7500, Test Acc: 0.7420\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Use a list of mixed modules and strings as the aggregation argument\n",
        "And the mix of them is supported as well for your convenience."
      ],
      "metadata": {
        "id": "Xf2Qb_ydL2w2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(seed)\n",
        "model = GNN(16, aggr=[\n",
        "    'mean',\n",
        "    MaxAggregation(),\n",
        "    'sum',\n",
        "    StdAggregation(),\n",
        "    'var',\n",
        "])\n",
        "print(model)\n",
        "run(model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0uLPbAHIHvYr",
        "outputId": "cf8adf13-6fca-46e7-8bd5-ce51a16e4abf"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GNN(\n",
            "  (conv1): SAGEConv(500, 16, aggr=['mean', 'MaxAggregation()', 'sum', 'StdAggregation()', 'var'])\n",
            "  (conv2): SAGEConv(16, 3, aggr=['mean', 'MaxAggregation()', 'sum', 'StdAggregation()', 'var'])\n",
            ")\n",
            "Epoch: 000, Train: 0.7833, Val Acc: 0.6420, Test Acc: 0.6320\n",
            "Epoch: 001, Train: 0.8333, Val Acc: 0.7040, Test Acc: 0.6980\n",
            "Epoch: 002, Train: 0.8500, Val Acc: 0.6660, Test Acc: 0.6570\n",
            "Epoch: 003, Train: 0.9500, Val Acc: 0.7120, Test Acc: 0.7010\n",
            "Epoch: 004, Train: 0.9333, Val Acc: 0.7500, Test Acc: 0.7420\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Define multiple aggregations with `MultiAggregation` module\n",
        "\n",
        "When a list is taken, [MessagePassing](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.MessagePassing) would stack these aggregators in via the [MultiAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.MultiAggregation) module automatically. But you can also directly pass a [MultiAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.MultiAggregation) instead of a list. Now let's see how can we define multiple aggregations with [MultiAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.MultiAggregation). Here we use different initial temperatures for [SoftmaxAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.SoftmaxAggregation) ([Li et al. (2020)](https://arxiv.org/abs/2006.07739)). Every different temperature will result in aggregation with different softness."
      ],
      "metadata": {
        "id": "dFw0MMx3L-rd"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(seed)\n",
        "aggr = MultiAggregation([\n",
        "    SoftmaxAggregation(t=0.01, learn=True),\n",
        "    SoftmaxAggregation(t=1, learn=True),\n",
        "    SoftmaxAggregation(t=100, learn=True),\n",
        "])\n",
        "model = GNN(16, aggr=aggr)\n",
        "print(model)\n",
        "run(model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D6iUX1IvICOW",
        "outputId": "1f4da92e-9b35-49fe-f020-e42dbb3757a2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GNN(\n",
            "  (conv1): SAGEConv(500, 16, aggr=MultiAggregation([\n",
            "    SoftmaxAggregation(learn=True),\n",
            "    SoftmaxAggregation(learn=True),\n",
            "    SoftmaxAggregation(learn=True)\n",
            "  ], mode=cat))\n",
            "  (conv2): SAGEConv(16, 3, aggr=MultiAggregation([\n",
            "    SoftmaxAggregation(learn=True),\n",
            "    SoftmaxAggregation(learn=True),\n",
            "    SoftmaxAggregation(learn=True)\n",
            "  ], mode=cat))\n",
            ")\n",
            "Epoch: 000, Train: 0.7333, Val Acc: 0.6080, Test Acc: 0.6340\n",
            "Epoch: 001, Train: 0.8500, Val Acc: 0.6440, Test Acc: 0.6560\n",
            "Epoch: 002, Train: 0.9333, Val Acc: 0.7420, Test Acc: 0.7390\n",
            "Epoch: 003, Train: 0.9333, Val Acc: 0.7560, Test Acc: 0.7360\n",
            "Epoch: 004, Train: 0.9833, Val Acc: 0.7840, Test Acc: 0.7460\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "What is more?\n",
        "There are many other aggregation operators supported for you to \"lego\" your GNNs. [PowerMeanAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.PowerMeanAggregation) allows you to define and potentially learn generalized means beyond simple  arithmetic mean such as harmonic mean and geometric mean. [LSTMAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.LSTMAggregation) can perform permutation-variant aggregation. More other interesting aggregation operators such as [Set2Set](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.Set2Set), [DegreeScalerAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.DegreeScalerAggregation), [SortAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.SortAggregation), [GraphMultisetTransformer](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.GraphMultisetTransformer), [AttentionalAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.AttentionalAggregation) and [EquilibriumAggregation](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.aggr.EquilibriumAggregation) are ready for you to explore."
      ],
      "metadata": {
        "id": "AHlCxX2LbLe2"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SDOmdUe0C3U1"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "In this tutorial, you have been presented with the `torch_geometric.nn.aggr` package which provides a flexible interface to experiment with different aggregation functions with your message passing convolutions and unifies aggregation within GNNs across [`MessagePassing`](https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/conv/message_passing.py) and [global readouts](https://github.com/pyg-team/pytorch_geometric/tree/master/torch_geometric/nn/glob). This new abstraction also makes designing new types of aggregation functions easier. Now, you can create your own aggregation function with the base `Aggregation` class. Please refer to the [docs](https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/nn/aggr/base.html#Aggregation) for more details.\n",
        "\n",
        "```python\n",
        "class MyAggregation(Aggregation):\n",
        "    def __init__(self, ...):\n",
        "      ...\n",
        "\n",
        "    def forward(self, x: Tensor, index: Optional[Tensor] = None,\n",
        "                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,\n",
        "                dim: int = -2) -> Tensor:\n",
        "      ...\n",
        "```\n",
        "\n",
        "*Have fun!*"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "aYk1EAZweU2y"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}