sgc.py

"""
This code was modified from implementations of SGC in other backends.

Simplifying Graph Convolutional Networks (Wu, Zhang and Souza et al, 2019)
Paper: https://arxiv.org/abs/1902.07153
Author Implementation: https://github.com/Tiiiger/SGC

SGC implementation in DGL.
"""
import argparse
import textwrap

import tensorflow as tf
import tensorflow_addons as tfa

from dgl.data import CiteseerGraphDataset, CoraGraphDataset, PubmedGraphDataset
from dgl.nn.tensorflow.conv import SGConv

_DATASETS = {
    "citeseer": CiteseerGraphDataset(verbose=False),
    "cora": CoraGraphDataset(verbose=False),
    "pubmed": PubmedGraphDataset(verbose=False),
}


def load_data(dataset):
    return _DATASETS[dataset]


def _sum_boolean_tensor(x):
    return tf.reduce_sum(tf.cast(x, dtype="int64"))


def describe_data(data):
    g = data[0]

    n_edges = g.number_of_edges()
    num_classes = data.num_classes

    train_mask = g.ndata["train_mask"]
    val_mask = g.ndata["val_mask"]
    test_mask = g.ndata["test_mask"]

    description = textwrap.dedent(
        f"""
        ----Data statistics----
        Edges           {n_edges:,.0f}
        Classes         {num_classes:,.0f}
        Train samples   {_sum_boolean_tensor(train_mask):,.0f}
        Val samples     {_sum_boolean_tensor(val_mask):,.0f}
        Test samples    {_sum_boolean_tensor(test_mask):,.0f}
        """
    )
    return description


class SGC(tf.keras.Model):
    def __init__(self, g, num_classes, bias=False):
        super().__init__()
        self.num_classes = num_classes
        self.g = self.ensure_self_loop(g)
        self.conv = SGConv(
            in_feats=self.in_feats,
            out_feats=self.num_classes,
            k=2,
            cached=True,
            bias=bias,
        )

    def call(self, inputs):
        return self.conv(self.g, inputs)

    @property
    def in_feats(self):
        return self.g.ndata["feat"].shape[1]

    @property
    def num_nodes(self):
        return self.g.num_nodes()

    @staticmethod
    def ensure_self_loop(g):
        g = g.remove_self_loop()
        g = g.add_self_loop()
        return g

    def train_step(self, data):
        X, y = data
        mask = self.g.ndata["train_mask"]

        with tf.GradientTape() as tape:
            y_pred = self(X, training=True)
            loss = self.compiled_loss(y[mask], y_pred[mask])

        trainable_variables = self.trainable_variables
        gradients = tape.gradient(loss, trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, trainable_variables))
        self.compiled_metrics.update_state(y[mask], y_pred[mask])
        return {m.name: m.result() for m in self.metrics}

    def test_step(self, data):
        X, y = data
        mask = self.g.ndata["val_mask"]
        y_pred = self(X, training=False)
        self.compiled_loss(y[mask], y_pred[mask])
        self.compiled_metrics.update_state(y[mask], y_pred[mask])
        return {m.name: m.result() for m in self.metrics}

    def compile(self, *args, **kwargs):
        super().compile(*args, **kwargs, run_eagerly=True)

    def fit(self, *args, **kwargs):
        kwargs["batch_size"] = self.num_nodes
        kwargs["shuffle"] = False
        super().fit(*args, **kwargs)

    def predict(self, *args, **kwargs):
        kwargs["batch_size"] = self.num_nodes
        return super().predict(*args, **kwargs)


def main(dataset, lr, bias, n_epochs, weight_decay):
    data = load_data(dataset)
    print(describe_data(data))

    g = data[0]
    X = g.ndata["feat"]
    y = g.ndata["label"]

    model = SGC(g=g, num_classes=data.num_classes, bias=bias)

    loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)
    optimizer = tfa.optimizers.AdamW(weight_decay, lr)
    accuracy = tf.metrics.SparseCategoricalAccuracy(name="accuracy")

    model.compile(optimizer, loss, metrics=[accuracy])
    model.fit(x=X, y=y, epochs=n_epochs, validation_data=(X, y))

    y_pred = model.predict(X, batch_size=len(X))
    test_mask = g.ndata["test_mask"]
    test_accuracy = accuracy(y[test_mask], y_pred[test_mask])
    print(f"Test Accuracy: {test_accuracy:.1%}")


def _parse_args():
    parser = argparse.ArgumentParser(
        description="Run experiment for Simple Graph Convolution (SGC)"
    )
    parser.add_argument("--dataset", default="cora", help="dataset to run")
    parser.add_argument("--lr", type=float, default=0.2, help="learning rate")
    parser.add_argument(
        "--bias", action="store_true", default=False, help="flag to use bias"
    )
    parser.add_argument(
        "--n-epochs", type=int, default=100, help="number of training epochs"
    )
    parser.add_argument(
        "--weight-decay", type=float, default=5e-6, help="weight for L2 loss"
    )
    return parser.parse_args()


if __name__ == "__main__":
    args = _parse_args()
    main(
        dataset=args.dataset,
        lr=args.lr,
        bias=args.bias,
        n_epochs=args.n_epochs,
        weight_decay=args.weight_decay,
    )