Source code for dgl.nn.tensorflow.conv.appnpconv

"""TF Module for APPNPConv"""
# pylint: disable= no-member, arguments-differ, invalid-name
import tensorflow as tf
from tensorflow.keras import layers
import numpy as np

from .... import function as fn


[docs]class APPNPConv(layers.Layer): r"""Approximate Personalized Propagation of Neural Predictions layer from `Predict then Propagate: Graph Neural Networks meet Personalized PageRank <https://arxiv.org/pdf/1810.05997.pdf>`__ .. math:: H^{0} & = X H^{t+1} & = (1-\alpha)\left(\hat{D}^{-1/2} \hat{A} \hat{D}^{-1/2} H^{t}\right) + \alpha H^{0} Parameters ---------- k : int Number of iterations :math:`K`. alpha : float The teleport probability :math:`\alpha`. edge_drop : float, optional Dropout rate on edges that controls the messages received by each node. Default: ``0``. """ def __init__(self, k, alpha, edge_drop=0.): super(APPNPConv, self).__init__() self._k = k self._alpha = alpha self.edge_drop = layers.Dropout(edge_drop) def call(self, graph, feat): r"""Compute APPNP layer. Parameters ---------- graph : DGLGraph The graph. feat : tf.Tensor The input feature of shape :math:`(N, *)` :math:`N` is the number of nodes, and :math:`*` could be of any shape. Returns ------- tf.Tensor The output feature of shape :math:`(N, *)` where :math:`*` should be the same as input shape. """ with graph.local_scope(): degs = tf.clip_by_value(tf.cast(graph.in_degrees(), tf.float32), clip_value_min=1, clip_value_max=np.inf) norm = tf.pow(degs, -0.5) shp = norm.shape + (1,) * (feat.ndim - 1) norm = tf.reshape(norm, shp) feat_0 = feat for _ in range(self._k): # normalization by src node feat = feat * norm graph.ndata['h'] = feat graph.edata['w'] = self.edge_drop( tf.ones(graph.number_of_edges(), 1)) graph.update_all(fn.u_mul_e('h', 'w', 'm'), fn.sum('m', 'h')) feat = graph.ndata.pop('h') # normalization by dst node feat = feat * norm feat = (1 - self._alpha) * feat + self._alpha * feat_0 return feat