Train a Hypergraph Message Passing Neural Network (HMPNN)

In this notebook, we will create and train a Hypergraph Message Passing Neural Network in the hypergraph domain. This method is introduced in the paper Message Passing Neural Networks for Hypergraphs by Heydari et Livi 2022. We will use a benchmark dataset, Cora, a collection of 2708 academic papers and 5429 citation relations, to do the task of node classification. There are 7 category labels, namely Case_Based, Genetic_Algorithms, Neural_Networks, Probabilistic_Methods, Reinforcement_Learning, Rule_Learning and Theory.

Each document is initially represented as a binary vector of length 1433, standing for a unique subset of the words within the papers, in which a value of 1 means the presence of its corresponding word in the paper.

import torch
import torch_geometric.datasets as geom_datasets
from sklearn.metrics import accuracy_score

from topomodelx.nn.hypergraph.hmpnn import HMPNN

torch.manual_seed(0)

If GPU’s are available, we will make use of them. Otherwise, this will run on CPU.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

cuda

Pre-processing

Here we download the dataset. It contains initial representation of nodes, the adjacency information, category labels and train-val-test masks.

dataset = geom_datasets.Planetoid(root="tmp/", name="cora")[0]

Below, we construct the incidence matrix (\(B_1\)) which is of shape \(n_\text{nodes} \times n_\text{edges}\).

dataset["incidence_1"] = torch.sparse_coo_tensor(
    dataset["edge_index"], torch.ones(dataset["edge_index"].shape[1]), dtype=torch.long
)
dataset = dataset.to(device)

x_0s = dataset["x"]
y = dataset["y"]
incidence_1 = dataset["incidence_1"]

Train the Neural Network

We then specify the hyperparameters and construct the model, the loss and optimizer.

class Network(torch.nn.Module):
    """Network class that initializes the base model and readout layer.

    Base model parameters:
    ----------
    Reqired:
    in_channels : int
        Dimension of the input features.
    hidden_channels : int
        Dimension of the hidden features.

    Optitional:
    **kwargs : dict
        Additional arguments for the base model.

    Readout layer parameters:
    ----------
    out_channels : int
        Dimension of the output features.
    task_level : str
        Level of the task. Either "graph" or "node".
    """

    def __init__(
        self, in_channels, hidden_channels, out_channels, task_level="graph", **kwargs
    ):
        super().__init__()

        # Define the model
        self.base_model = HMPNN(
            in_channels=in_channels, hidden_channels=hidden_channels, **kwargs
        )

        # Readout
        self.linear = torch.nn.Linear(hidden_channels, out_channels)
        self.out_pool = task_level == "graph"

    def forward(self, x_0, x_1, incidence_1):
        # Base model
        x_0, x_1 = self.base_model(x_0, x_1, incidence_1)

        # Pool over all nodes in the hypergraph
        x = torch.max(x_0, dim=0)[0] if self.out_pool is True else x_0

        return self.linear(x)

Initialize the model

# Base model hyperparameters
in_channels = x_0s.shape[1]
hidden_channels = 128
n_layers = 1

# Readout hyperparameters
out_channels = torch.unique(y).shape[0]
task_level = "graph" if out_channels == 1 else "node"


model = Network(
    in_channels=in_channels,
    hidden_channels=hidden_channels,
    out_channels=out_channels,
    n_layers=n_layers,
    task_level=task_level,
).to(device)

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
loss_fn = torch.nn.CrossEntropyLoss()


train_mask = dataset["train_mask"]
val_mask = dataset["val_mask"]
test_mask = dataset["test_mask"]

Now it’s time to train the model, looping over the network for a low amount of epochs. We keep training minimal for the purpose of rapid testing.

Note: The number of epochs below have been kept low to facilitate debugging and testing. Real use cases should likely require more epochs.

torch.manual_seed(0)
test_interval = 5
num_epochs = 5


initial_x_1 = torch.zeros_like(x_0s)
for epoch in range(1, num_epochs + 1):
    model.train()
    optimizer.zero_grad()
    y_hat = model(x_0s, initial_x_1, incidence_1)
    loss = loss_fn(y_hat[train_mask], y[train_mask])
    loss.backward()
    optimizer.step()

    train_loss = loss.item()
    y_pred = y_hat.argmax(dim=-1)
    train_acc = accuracy_score(y[train_mask].cpu(), y_pred[train_mask].cpu())

    if epoch % test_interval == 0:
        model.eval()

        y_hat = model(x_0s, initial_x_1, incidence_1)
        val_loss = loss_fn(y_hat[val_mask], y[val_mask]).item()
        y_pred = y_hat.argmax(dim=-1)
        val_acc = accuracy_score(y[val_mask].cpu(), y_pred[val_mask].cpu())

        test_loss = loss_fn(y_hat[test_mask], y[test_mask]).item()
        y_pred = y_hat.argmax(dim=-1)
        test_acc = accuracy_score(y[test_mask].cpu(), y_pred[test_mask].cpu())
        print(
            f"Epoch: {epoch + 1} train loss: {train_loss:.4f} train acc: {train_acc:.2f} "
            f" val loss: {val_loss:.4f} val acc: {val_acc:.2f}"
            f" test loss: {test_acc:.4f} val acc: {test_acc:.2f}"
        )

Epoch: 6 train loss: 1.2665 train acc: 0.82  val loss: 1.9848 val acc: 0.23 test loss: 0.2270 val acc: 0.23
Epoch: 11 train loss: 0.8439 train acc: 0.99  val loss: 1.7569 val acc: 0.38 test loss: 0.3970 val acc: 0.40
Epoch: 16 train loss: 0.5119 train acc: 1.00  val loss: 1.6846 val acc: 0.39 test loss: 0.4160 val acc: 0.42
Epoch: 21 train loss: 0.2717 train acc: 1.00  val loss: 1.5872 val acc: 0.43 test loss: 0.4500 val acc: 0.45
Epoch: 26 train loss: 0.1571 train acc: 1.00  val loss: 1.6143 val acc: 0.41 test loss: 0.4230 val acc: 0.42
Epoch: 31 train loss: 0.0816 train acc: 1.00  val loss: 1.5894 val acc: 0.45 test loss: 0.4490 val acc: 0.45
Epoch: 36 train loss: 0.0478 train acc: 1.00  val loss: 1.6020 val acc: 0.46 test loss: 0.4630 val acc: 0.46
Epoch: 41 train loss: 0.0298 train acc: 1.00  val loss: 1.6153 val acc: 0.47 test loss: 0.4670 val acc: 0.47
Epoch: 46 train loss: 0.0214 train acc: 1.00  val loss: 1.6499 val acc: 0.47 test loss: 0.4720 val acc: 0.47
Epoch: 51 train loss: 0.0160 train acc: 1.00  val loss: 1.6764 val acc: 0.48 test loss: 0.4830 val acc: 0.48
Epoch: 56 train loss: 0.0149 train acc: 1.00  val loss: 1.6986 val acc: 0.48 test loss: 0.4900 val acc: 0.49
Epoch: 61 train loss: 0.0123 train acc: 1.00  val loss: 1.6888 val acc: 0.47 test loss: 0.4920 val acc: 0.49
Epoch: 66 train loss: 0.0097 train acc: 1.00  val loss: 1.6670 val acc: 0.48 test loss: 0.4970 val acc: 0.50
Epoch: 71 train loss: 0.0078 train acc: 1.00  val loss: 1.6547 val acc: 0.49 test loss: 0.5030 val acc: 0.50
Epoch: 76 train loss: 0.0072 train acc: 1.00  val loss: 1.6484 val acc: 0.49 test loss: 0.5030 val acc: 0.50
Epoch: 81 train loss: 0.0066 train acc: 1.00  val loss: 1.6378 val acc: 0.49 test loss: 0.5100 val acc: 0.51
Epoch: 86 train loss: 0.0064 train acc: 1.00  val loss: 1.6507 val acc: 0.49 test loss: 0.5110 val acc: 0.51
Epoch: 91 train loss: 0.0060 train acc: 1.00  val loss: 1.6745 val acc: 0.50 test loss: 0.5100 val acc: 0.51
Epoch: 96 train loss: 0.0051 train acc: 1.00  val loss: 1.6682 val acc: 0.50 test loss: 0.5150 val acc: 0.52
Epoch: 101 train loss: 0.0047 train acc: 1.00  val loss: 1.6412 val acc: 0.50 test loss: 0.5190 val acc: 0.52