Very deep GNNs for protein function prediction (DeeperGCN)
Author: K3-Node Team
Backend: Multi-Backend
Dataset: PPI
Description: DeeperGCN (Li et al., 2020) trains GNNs with 28 layers using pre-activation residual blocks (DeepGCNLayer, "res+") and GENConv layers with a learnable softmax aggregation.
Very deep GNNs for protein function prediction (DeeperGCN)
DeeperGCN (Li et al., 2020) trains GNNs with 28 layers using
pre-activation residual blocks (DeepGCNLayer, "res+") and GENConv layers with a learnable softmax
aggregation. The task is multi-label: predict the functions of proteins in protein interaction
graphs.
Same model as PyG's examples/ogbn_proteins_deepgcn.py; on PPI instead of ogbn-proteins, trained for 200 epochs instead of PyG's 1,000 to keep the notebook quick.
Install K3-Node, then choose a backend: "tensorflow", "torch" or "jax"
Load the data
PyG's example uses ogbn-proteins (132,534 proteins, 112 functions). Here the 20 training graphs of PPI (121 functions) are combined into one graph and split into 10 random parts per epoch. PPI has no edge features, so every edge gets the same one.
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.layers import GENConv
from k3_node.loader import DataLoader, RandomNodeLoader
from k3_node.metrics import F1Score
from k3_node.models import DeepGCNLayer
def add_edge_features(data):
data.edge_attr = np.ones((data.edge_index.shape[1], 1), dtype="float32")
return data
train_dataset = PPI("data/PPI", split="train", transform=add_edge_features)
test_dataset = PPI("data/PPI", split="test", transform=add_edge_features)
train_data = Batch.from_data_list(list(train_dataset))
train_loader = RandomNodeLoader(train_data, num_parts=10, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=2)
Define the model
class DeeperGCN(keras.Model):
def __init__(self, hidden_channels, num_layers, out_channels):
super().__init__()
self.node_encoder = keras.layers.Dense(hidden_channels)
self.edge_encoder = keras.layers.Dense(hidden_channels)
self.layers_ = []
for _ in range(num_layers):
conv = GENConv(hidden_channels, hidden_channels, aggr="softmax", t=1.0, learn_t=True, num_layers=2,
norm="layer")
self.layers_.append(DeepGCNLayer(conv, keras.layers.LayerNormalization(), keras.layers.ReLU(),
block="res+", dropout=0.1))
self.dropout = keras.layers.Dropout(0.1)
self.lin = keras.layers.Dense(out_channels)
def call(self, data, training=False):
x = self.node_encoder(data.x)
edge_attr = self.edge_encoder(data.edge_attr)
x = self.layers_[0].conv(x, data.edge_index, edge_attr)
for layer in self.layers_[1:]:
x = layer(x, data.edge_index, edge_attr, training=training)
x = self.layers_[0].act(self.layers_[0].norm(x))
return self.lin(self.dropout(x, training=training))
model = DeeperGCN(hidden_channels=64, num_layers=28, out_channels=train_dataset.num_classes)
Train
Each of the 121 outputs is a yes/no prediction (binary cross-entropy); PyG reports the ROC-AUC.
model.compile(
optimizer=keras.optimizers.Adam(learning_rate=0.01),
loss=keras.losses.BinaryCrossentropy(from_logits=True),
metrics=[F1Score(average="micro", from_logits=True)],
)
model.fit(train_loader, epochs=200, verbose=2)
Evaluate
Like OGB's evaluator, compute the ROC-AUC of every function and average them.