layers.py

import tensorflow as tf
from flags import *
from inits import *

flags = tf.compat.v1.flags
FLAGS = flags.FLAGS

# global unique layer ID dictionary for layer name assignment
_LAYER_UIDS = {}


def get_layer_uid(layer_name=''):
  """Helper function, assigns unique layer IDs."""
  if layer_name not in _LAYER_UIDS:
    _LAYER_UIDS[layer_name] = 1
    return 1
  else:
    _LAYER_UIDS[layer_name] += 1
    return _LAYER_UIDS[layer_name]


def sparse_dropout(x, keep_prob, noise_shape):
  """Dropout for sparse tensors."""
  random_tensor = keep_prob
  random_tensor += tf.random.uniform(noise_shape)
  dropout_mask = tf.cast(tf.floor(random_tensor), dtype=tf.bool)
  pre_out = tf.sparse_retain(x, dropout_mask)
  return pre_out * (1. / keep_prob)


def dot(x, y, sparse=False):
  """Wrapper for tf.matmul (sparse vs dense)."""
  if sparse:
    res = tf.sparse_tensor_dense_matmul(x, y)
  else:
    res = tf.matmul(x, y)
  return res


class Layer(object):
  """Base layer class. Defines basic API for all layer objects.
    Implementation inspired by keras (http://keras.io).

    # Properties
        name: String, defines the variable scope of the layer.
        logging: Boolean, switches Tensorflow histogram logging on/off

    # Methods
        _call(inputs): Defines computation graph of layer
            (i.e. takes input, returns output)
        __call__(inputs): Wrapper for _call()
        _log_vars(): Log all variables
    """

  def __init__(self, **kwargs):
    allowed_kwargs = {'name', 'logging'}
    for kwarg in kwargs.keys():
      assert kwarg in allowed_kwargs, 'Invalid keyword argument: ' + kwarg
    name = kwargs.get('name')
    if not name:
      layer = self.__class__.__name__.lower()
      name = layer + '_' + str(get_layer_uid(layer))
    self.name = name
    self.vars = {}
    logging = kwargs.get('logging', False)
    self.logging = logging
    self.sparse_inputs = False

  def _call(self, inputs):
    return inputs

  def __call__(self, inputs):
    with tf.name_scope(self.name):
      if self.logging and not self.sparse_inputs:
        tf.summary.histogram(self.name + '/inputs', inputs)
      outputs = self._call(inputs)
      if self.logging:
        tf.summary.histogram(self.name + '/outputs', outputs)
      return outputs

  def _log_vars(self):
    for var in self.vars:
      tf.summary.histogram(self.name + '/vars/' + var, self.vars[var])


class CustomDense(Layer):
  """Dense layer."""

  def __init__(self,
               input_dim,
               output_dim,
               name=None,
               dropout=FLAGS.use_dropout,
               dropout_rate=FLAGS.vae_dropout_rate,
               sparse_inputs=False,
               act=tf.nn.relu,
               bias=True,
               **kwargs):
    super(CustomDense, self).__init__(**kwargs)
    if name is None:
      name = self.name
    if dropout:
      self.dropout = dropout_rate
    else:
      self.dropout = 0.

    self.act = act
    self.sparse_inputs = sparse_inputs
    self.bias = bias

    with tf.compat.v1.variable_scope(name + '_vars'):
      self.vars['weights'] = glorot([input_dim, output_dim], name='weights')
      if self.bias:
        self.vars['bias'] = zeros([output_dim], name='bias')

    if self.logging:
      self._log_vars()

  def _call(self, inputs):
    x = inputs
    # dropout
    if self.sparse_inputs:
      x = sparse_dropout(x, 1 - self.dropout, self.num_features_nonzero)
    else:
      x = tf.nn.dropout(x, 1 - self.dropout)

    # transform
    output = dot(x, self.vars['weights'], sparse=self.sparse_inputs)

    # bias
    if self.bias:
      output += self.vars['bias']

    return self.act(output)


class GraphConvolution(Layer):
  """Basic graph convolution layer for undirected graph without edge labels."""

  def __init__(self,
               input_dim,
               output_dim,
               adj,
               dropout=FLAGS.use_dropout,
               dropout_rate=FLAGS.vae_dropout_rate,
               act=tf.nn.relu,
               **kwargs):
    super(GraphConvolution, self).__init__(**kwargs)
    with tf.compat.v1.variable_scope(self.name + '_vars'):
      self.vars['weights'] = glorot([input_dim, output_dim], name="weights")
    self.dropout = dropout_rate
    self.adj = adj
    self.act = act

  def _call(self, inputs):
    x = inputs
    x = tf.nn.dropout(x, rate=self.dropout)
    x = tf.matmul(x, self.vars['weights'])
    x = tf.sparse.sparse_dense_matmul(self.adj, x)
    outputs = self.act(x)
    return outputs


class GraphConvolutionSparse(Layer):
  """Graph convolution layer for sparse inputs."""

  def __init__(self,
               input_dim,
               output_dim,
               adj,
               features_nonzero,
               dropout=FLAGS.use_dropout,
               dropout_rate=FLAGS.vae_dropout_rate,
               act=tf.nn.relu,
               **kwargs):
    super(GraphConvolutionSparse, self).__init__(**kwargs)
    with tf.compat.v1.variable_scope(self.name + '_vars'):
      self.vars['weights'] = glorot([input_dim, output_dim], name="weights")
    self.dropout = dropout_rate
    self.adj = adj
    self.act = act
    self.issparse = True
    self.features_nonzero = features_nonzero

  def _call(self, inputs):
    x = inputs
    x = dropout_sparse(x, 1 - self.dropout, self.features_nonzero)
    x = tf.sparse_tensor_dense_matmul(x, self.vars['weights'])
    x = tf.sparse_tensor_dense_matmul(self.adj, x)
    outputs = self.act(x)
    return outputs


class InnerProductDecoder(Layer):
  """Decoder model layer for link prediction."""

  def __init__(self,
               input_dim,
               dropout=FLAGS.use_dropout,
               dropout_rate=FLAGS.vae_dropout_rate,
               act=tf.nn.sigmoid,
               **kwargs):
    super(InnerProductDecoder, self).__init__(**kwargs)
    self.dropout = dropout_rate
    self.act = act

  def _call(self, inputs):
    inputs = tf.nn.dropout(inputs, 1 - self.dropout)
    x = tf.transpose(inputs)
    x = tf.matmul(inputs, x)
    #x = tf.reshape(x, [-1])
    outputs = self.act(x)
    return outputs