Unsupervised GraphSAGE on protein graphs (PPI)

Author: K3-Node Team
Backend: Multi-Backend
Dataset: PPI
Description: Learn embeddings of the proteins in the PPI graphs without labels: GraphSAGE (Hamilton et al., 2017) is trained to predict which proteins interact.

View in Colab   GitHub source

---

Unsupervised GraphSAGE on protein graphs (PPI)

Learn embeddings of the proteins in the PPI graphs without labels: GraphSAGE (Hamilton et al., 2017) is trained to predict which proteins interact. A linear classifier on the embeddings then predicts each protein's 121 functions on graphs the model has never seen.

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

All training graphs are combined into one; LinkNeighborLoader yields batches of 2,048 edges with their sampled neighborhoods, plus as many random non-edges.

import keras
import numpy as np
from keras import ops
from k3_node.data import Batch
from k3_node.datasets import PPI
from k3_node.loader import DataLoader, LinkNeighborLoader
from k3_node.models import GraphSAGE

train_dataset = PPI("data/PPI", split="train")
val_dataset = PPI("data/PPI", split="val")
test_dataset = PPI("data/PPI", split="test")
train_data = Batch.from_data_list(list(train_dataset))
loader = LinkNeighborLoader(train_data, batch_size=2048, 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, out_channels=64)

    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(train_dataset.num_features)

Train

model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.005),
              loss=keras.losses.BinaryCrossentropy(from_logits=True))
model.fit(loader, epochs=5, verbose=2)

Evaluate

A linear classifier (one per function) is trained on the embeddings of the training proteins and scored with the micro-averaged F1 score.

from sklearn.linear_model import SGDClassifier
from sklearn.metrics import f1_score
from sklearn.multioutput import MultiOutputClassifier


def encode(dataset):
    loader = DataLoader(dataset, batch_size=2)
    z = model.predict(loader, verbose=0)
    y = np.concatenate([ops.convert_to_numpy(graph.y) for graph in dataset])
    return z, y


z, y = encode(train_dataset)
clf = MultiOutputClassifier(SGDClassifier(loss="log_loss", penalty="l2")).fit(z, y)
for name, split in [("Validation", val_dataset), ("Test", test_dataset)]:
    z, y = encode(split)
    print(f"{name} F1: {f1_score(y, clf.predict(z), average='micro'):.4f}")