Skip to content

SuperGAT on Cora

Author: K3-Node Team
Backend: Multi-Backend
Dataset: Cora (Planetoid)
Description: Classify papers in the Cora citation network by topic.

View in Colab   GitHub source

SuperGAT on Cora

Classify papers in the Cora citation network by topic. SuperGAT (Kim & Oh, 2021) is a graph attention network whose attention is also trained on a self-supervised task: telling real links apart from random pairs of papers. attention_loss_weight=4.0 adds that task's loss to the model loss during training.

Same model as PyG's examples/super_gat.py; trained for 200 epochs instead of PyG's 500 to keep the notebook quick.

Install K3-Node, then choose a backend: "tensorflow", "torch" or "jax"

!pip install k3-node[examples]
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

Load the data

data holds one graph: node features x, edges edge_index, labels y, and masks marking the training, validation and test nodes.

import keras
from keras import ops
from k3_node.datasets import Planetoid
from k3_node.loader import FullGraphDataset
from k3_node.transforms import NormalizeFeatures

dataset = Planetoid("data/Planetoid", name="Cora", transform=NormalizeFeatures())
data = dataset[0]
print(data)

Define the model

from k3_node.layers import SuperGATConv


class SuperGAT(keras.Model):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        options = dict(heads=8, dropout=0.6, attention_type="MX", edge_sample_ratio=0.8,
                       is_undirected=True, attention_loss_weight=4.0)
        self.dropout = keras.layers.Dropout(0.6)
        self.conv1 = SuperGATConv(in_channels, 8, **options)
        self.conv2 = SuperGATConv(8 * 8, out_channels, concat=False, **options)

    def call(self, data, training=False):
        x = self.dropout(data.x, training=training)
        x = ops.elu(self.conv1(x, data.edge_index, training=training))
        x = self.dropout(x, training=training)
        return self.conv2(x, data.edge_index, training=training)


model = SuperGAT(dataset.num_features, dataset.num_classes)

Train

FullGraphDataset feeds the whole graph to Keras; mask selects which nodes count in the loss and the accuracy.

model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.005, weight_decay=0.0005),
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    weighted_metrics=["accuracy"],
)
model.fit(
    FullGraphDataset(data, mask="train_mask"),
    validation_data=FullGraphDataset(data, mask="val_mask"),
    epochs=200,
    verbose=2,
)

Evaluate

loss, accuracy = model.evaluate(FullGraphDataset(data, mask="test_mask"), verbose=0)
print(f"Test accuracy: {accuracy:.4f}")