Skip to content

Unsupervised GraphSAGE (Cora)

Author: K3-Node Team
Backend: Multi-Backend
Dataset: Cora (Planetoid)
Description: Learn node embeddings of the Cora graph without labels: GraphSAGE (Hamilton et al., 2017) is trained to predict which pairs of nodes are linked.

View in Colab   GitHub source

Unsupervised GraphSAGE (Cora)

Learn node embeddings of the Cora graph without labels: GraphSAGE (Hamilton et al., 2017) is trained to predict which pairs of nodes are linked. The embeddings are then evaluated by how well a simple classifier predicts the topics of the papers.

Same model as PyG's examples/graph_sage_unsup.py.

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

LinkNeighborLoader yields mini-batches of edges together with the sampled neighborhoods of their endpoints.

import keras
from keras import ops
from k3_node.datasets import Planetoid
from k3_node.loader import FullGraphDataset, LinkNeighborLoader
from k3_node.models import GraphSAGE
from k3_node.transforms import NormalizeFeatures

dataset = Planetoid("data/Planetoid", name="Cora", transform=NormalizeFeatures())
data = dataset[0]
train_loader = LinkNeighborLoader(data, batch_size=256, shuffle=True, neg_sampling_ratio=1.0, num_neighbors=[10, 10])

Define the model

class LinkModel(keras.Model):
    def __init__(self, in_channels):
        super().__init__()
        self.sage = GraphSAGE(in_channels, hidden_channels=64, num_layers=2)

    def call(self, data, training=False):
        h = self.sage(data.x, data.edge_index, training=training)
        if getattr(data, "edge_label_index", None) is None:
            return h  # node embeddings
        src, dst = data.edge_label_index[0], data.edge_label_index[1]
        return ops.sum(ops.take(h, src, axis=0) * ops.take(h, dst, axis=0), axis=-1)


model = LinkModel(dataset.num_features)

Train

Each batch is a set of 256 edges (with their sampled neighborhoods) plus as many random non-edges; the model learns to score true edges high. No labels are used.

model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.01),
              loss=keras.losses.BinaryCrossentropy(from_logits=True))
model.fit(train_loader, epochs=50, verbose=2)

Evaluate

from sklearn.linear_model import LogisticRegression

z = model.predict(FullGraphDataset(data), verbose=0)
train_mask, val_mask, test_mask, y = (ops.convert_to_numpy(v) for v in (data.train_mask, data.val_mask,
                                                                        data.test_mask, data.y))
clf = LogisticRegression().fit(z[train_mask], y[train_mask])
print(f"Validation accuracy: {clf.score(z[val_mask], y[val_mask]):.4f}")
print(f"Test accuracy: {clf.score(z[test_mask], y[test_mask]):.4f}")