dgi.py

"""
Deep Graph Infomax in DGL

References
----------
Papers: https://arxiv.org/abs/1809.10341
Author's code: https://github.com/PetarV-/DGI
"""

import tensorflow as tf
from tensorflow.keras import layers
import numpy as np
import math
from gcn import GCN


class Encoder(layers.Layer):
    def __init__(self, g, in_feats, n_hidden, n_layers, activation, dropout):
        super(Encoder, self).__init__()
        self.g = g
        self.conv = GCN(g, in_feats, n_hidden, n_hidden,
                        n_layers, activation, dropout)

    def call(self, features, corrupt=False):
        if corrupt:
            perm = np.random.permutation(self.g.number_of_nodes())
            features = tf.gather(features, perm)
        features = self.conv(features)
        return features


class Discriminator(layers.Layer):
    def __init__(self, n_hidden):
        super(Discriminator, self).__init__()
        uinit = tf.keras.initializers.RandomUniform(
            -1.0/math.sqrt(n_hidden), 1.0/math.sqrt(n_hidden))
        self.weight = tf.Variable(initial_value=uinit(
            shape=(n_hidden, n_hidden), dtype='float32'), trainable=True)

    def call(self, features, summary):
        features = tf.matmul(features, tf.matmul(
            self.weight, tf.expand_dims(summary, -1)))
        return features


class DGI(tf.keras.Model):
    def __init__(self, g, in_feats, n_hidden, n_layers, activation, dropout):
        super(DGI, self).__init__()
        self.encoder = Encoder(g, in_feats, n_hidden,
                               n_layers, activation, dropout)
        self.discriminator = Discriminator(n_hidden)
        self.loss = tf.nn.sigmoid_cross_entropy_with_logits

    def call(self, features):
        positive = self.encoder(features, corrupt=False)
        negative = self.encoder(features, corrupt=True)
        summary = tf.nn.sigmoid(tf.reduce_mean(positive, axis=0))

        positive = self.discriminator(positive, summary)
        negative = self.discriminator(negative, summary)

        l1 = self.loss(tf.ones(positive.shape),positive)
        l2 = self.loss(tf.zeros(negative.shape), negative)

        return tf.reduce_mean(l1) + tf.reduce_mean(l2)


class Classifier(layers.Layer):
    def __init__(self, n_hidden, n_classes):
        super(Classifier, self).__init__()
        self.fc = layers.Dense(n_classes)

    def call(self, features):
        features = self.fc(features)
        return features