""" 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