Merge pull request #433 from JonasBergHansen/develop

Implementation of Total Variation Graph Neural Networks

.github/workflows/examples.yml

@@ -9,10 +9,10 @@ jobs:
 
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.6
+    - name: Set up Python 3.11
       uses: actions/setup-python@v2
       with:
-        python-version: 3.6
+        python-version: 3.11
     - name: Install dependencies
       run: |
         pip install ogb matplotlib

.github/workflows/style_check.yml

@@ -9,10 +9,10 @@ jobs:
 
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.6
+    - name: Set up Python 3.11
       uses: actions/setup-python@v2
       with:
-        python-version: 3.6
+        python-version: 3.11
     - name: Lint Python code
       run: |
         pip install flake8

.github/workflows/test.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        python-version: [3.7, 3.8, 3.9]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
         os: [ubuntu-latest, macos-latest, windows-latest]
 
     steps:

docs/autogen.py

@@ -40,6 +40,7 @@
             layers.GCSConv,
             layers.GINConv,
             layers.GraphSageConv,
+            layers.GTVConv,
             layers.TAGConv,
             layers.XENetConv,
         ],
@@ -50,6 +51,7 @@
         "methods": [],
         "classes": [
             layers.SRCPool,
+            layers.AsymCheegerCutPool,
             layers.DiffPool,
             layers.LaPool,
             layers.MinCutPool,

examples/graph_prediction/qm9_ecc.py

@@ -40,6 +40,7 @@
 loader_tr = DisjointLoader(dataset_tr, batch_size=batch_size, epochs=epochs)
 loader_te = DisjointLoader(dataset_te, batch_size=batch_size, epochs=1)
 
+
 ################################################################################
 # Build model
 ################################################################################

examples/other/node_clustering_tvgnn.py

@@ -0,0 +1,135 @@
+"""
+This example implements the node clustering experiment on citation networks
+from the paper:
+
+Total Variation Graph Neural Networks (https://arxiv.org/abs/2211.06218)
+Jonas Berg Hansen and Filippo Maria Bianchi
+"""
+
+import numpy as np
+import tensorflow as tf
+from sklearn.metrics.cluster import (
+    completeness_score,
+    homogeneity_score,
+    normalized_mutual_info_score,
+)
+from tensorflow.keras import Model
+from tqdm import tqdm
+
+from spektral.datasets import DBLP
+from spektral.datasets.citation import Citation
+from spektral.layers import AsymCheegerCutPool, GTVConv
+from spektral.utils.sparse import sp_matrix_to_sp_tensor
+
+tf.random.set_seed(1)
+
+################################
+# CONFIG/HYPERPARAMETERS
+################################
+dataset_id = "cora"
+mp_channels = 512
+mp_layers = 2
+mp_activation = "elu"
+delta_coeff = 0.311
+epsilon = 1e-3
+mlp_hidden_channels = 256
+mlp_hidden_layers = 1
+mlp_activation = "relu"
+totvar_coeff = 0.785
+balance_coeff = 0.514
+learning_rate = 1e-3
+epochs = 500
+
+################################
+# LOAD DATASET
+################################
+if dataset_id in ["cora", "citeseer", "pubmed"]:
+    dataset = Citation(dataset_id, normalize_x=True)
+elif dataset_id == "dblp":
+    dataset = DBLP(normalize_x=True)
+X = dataset.graphs[0].x
+A = dataset.graphs[0].a
+Y = dataset.graphs[0].y
+y = np.argmax(Y, axis=-1)
+n_clust = Y.shape[-1]
+
+
+################################
+# MODEL
+################################
+class ClusteringModel(Model):
+    """
+    Defines the general model structure
+    """
+
+    def __init__(self, aggr, pool):
+        super().__init__()
+
+        self.mp = aggr
+        self.pool = pool
+
+    def call(self, inputs):
+        x, a = inputs
+
+        out = x
+        for _mp in self.mp:
+            out = _mp([out, a])
+
+        _, _, s_pool = self.pool([out, a])
+
+        return s_pool
+
+
+# Define the message-passing layers
+MP_layers = [
+    GTVConv(
+        mp_channels, delta_coeff=delta_coeff, epsilon=1e-3, activation=mp_activation
+    )
+    for _ in range(mp_layers)
+]
+
+# Define the pooling layer
+pool_layer = AsymCheegerCutPool(
+    n_clust,
+    mlp_hidden=[mlp_hidden_channels for _ in range(mlp_hidden_layers)],
+    mlp_activation=mlp_activation,
+    totvar_coeff=totvar_coeff,
+    balance_coeff=balance_coeff,
+    return_selection=True,
+)
+
+# Instantiate model and optimizer
+model = ClusteringModel(aggr=MP_layers, pool=pool_layer)
+opt = tf.keras.optimizers.Adam(learning_rate=learning_rate)
+
+
+################################
+# TRAINING
+################################
+@tf.function(input_signature=None)
+def train_step(model, inputs):
+    with tf.GradientTape() as tape:
+        _ = model(inputs, training=True)
+        loss = sum(model.losses)
+    gradients = tape.gradient(loss, model.trainable_variables)
+    opt.apply_gradients(zip(gradients, model.trainable_variables))
+    return model.losses
+
+
+A = sp_matrix_to_sp_tensor(A)
+inputs = [X, A]
+loss_history = []
+
+# Training loop
+for _ in tqdm(range(epochs)):
+    outs = train_step(model, inputs)
+
+################################
+# INFERENCE/RESULTS
+################################
+S_ = model(inputs, training=False)
+s_out = np.argmax(S_, axis=-1)
+nmi = normalized_mutual_info_score(y, s_out)
+hom = homogeneity_score(y, s_out)
+com = completeness_score(y, s_out)
+print("Homogeneity: {:.3f}; Completeness: {:.3f}; NMI: {:.3f}".format(hom, com, nmi))

spektral/layers/base.py

@@ -118,7 +118,6 @@ def __init__(
         kernel_constraint=None,
         **kwargs
     ):
-
         super().__init__(**kwargs)
         self.trainable_kernel = trainable_kernel
         self.activation = activations.get(activation)
@@ -184,7 +183,6 @@ class MinkowskiProduct(Layer):
     """
 
     def __init__(self, activation=None, **kwargs):
-
         super().__init__(**kwargs)
         self.activation = activations.get(activation)
 

spektral/layers/convolutional/__init__.py

@@ -14,6 +14,7 @@
 from .general_conv import GeneralConv
 from .gin_conv import GINConv
 from .graphsage_conv import GraphSageConv
+from .gtv_conv import GTVConv
 from .message_passing import MessagePassing
 from .tag_conv import TAGConv
 from .xenet_conv import XENetConv, XENetDenseConv