Skip to content

Semi-supervised node classification with OGC

Author: K3-Node Team
Backend: Multi-Backend
Dataset: Cora (Planetoid)
Description: OGC (Wang et al., 2023) classifies the nodes of Cora with a linear classifier W and no deep network.

View in Colab   GitHub source

Semi-supervised node classification with OGC

OGC (Wang et al., 2023) classifies the nodes of Cora with a linear classifier W and no deep network. It alternates two steps: train W on the current node features U, then update U itself, smoothing it over the graph (lazy graph convolution) and nudging the training nodes towards their labels. It stops when the predictions stop changing.

Same model as PyG's examples/ogc.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

GCNNorm adds the normalized edge weights of a GCN. The lazy adjacency matrix keeps 90% of each node's features and mixes in 10% from its neighbors.

import keras
from keras import ops
from k3_node.datasets import Planetoid
from k3_node.layers import to_dense_adj
from k3_node.training import gradient_step
from k3_node.transforms import Compose, GCNNorm, NormalizeFeatures

decline = 0.9        # decay of the supervised step size
eta_sup = 0.001      # step size of the supervised update of U
eta_W = 0.5          # learning rate of W
beta = 0.1           # how much a node mixes in from its neighbors
max_sim_tol = 0.995  # predictions count as unchanged above this agreement
max_patience = 2     # stop after this many unchanged iterations

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

adj = to_dense_adj(data.edge_index, edge_attr=data.edge_weight, max_num_nodes=data.num_nodes)[0]
lazy_adj = beta * adj + (1 - beta) * ops.eye(data.num_nodes)
y_one_hot = ops.one_hot(data.y, dataset.num_classes)
train_mask = ops.cast(data.train_mask, "float32")[:, None]
trainval_mask = ops.cast(ops.logical_or(data.train_mask, data.val_mask), "float32")[:, None]

The two update steps

W = keras.layers.Dense(dataset.num_classes, use_bias=False)
optimizer = keras.optimizers.SGD(learning_rate=eta_W)


def update_W(U):  # one gradient step of least squares on the training and validation nodes
    gradient_step(lambda: ops.sum(trainval_mask * ops.square(W(U) - y_one_hot)), W.trainable_variables, optimizer)
    return W(U)


def update_U(U, pred):  # smooth over the graph, then correct towards the training labels
    U = ops.matmul(lazy_adj, U) - eta_sup * 2 * ops.matmul(train_mask * (pred - y_one_hot), ops.transpose(W.kernel))
    return ops.stop_gradient(U)  # U is data for the next step of W, not something to differentiate


def accuracy(pred, mask):
    correct = ops.cast(ops.argmax(pred, axis=-1) == ops.cast(data.y, "int32"), "float32")
    return float(ops.sum(correct * mask[:, 0]) / ops.sum(mask))

Run OGC

U = data.x
last_pred = ops.argmax(W(U), axis=-1)
patience = 0
for i in range(1, 65):
    U = update_U(U, update_W(U))
    eta_sup *= decline
    pred = W(U)
    test_accuracy = accuracy(pred, ops.cast(data.test_mask, "float32")[:, None])
    print(f"Iteration {i:02d}: train+val accuracy {accuracy(pred, trainval_mask):.4f}, test accuracy {test_accuracy:.4f}")

    labels = ops.argmax(pred, axis=-1)
    if float(ops.mean(ops.cast(labels == last_pred, "float32"))) > max_sim_tol:
        patience += 1
        if patience > max_patience:
            break
    last_pred = labels
print(f"Test accuracy: {test_accuracy:.4f}")