wehub/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dgl/nn/tensorflow/__init__.py
+6
View File
@@ -0,0 +1,6 @@
"""Package for Tensorflow-specific NN modules."""
from .conv import *
from .glob import *
from .hetero import *
from .softmax import *
from .utils import *
python/dgl/nn/tensorflow/conv/__init__.py
+11
View File
@@ -0,0 +1,11 @@
"""TF NN conv module"""
from .appnpconv import APPNPConv
from .chebconv import ChebConv
from .densechebconv import DenseChebConv
from .edgeconv import EdgeConv
from .gatconv import GATConv
from .ginconv import GINConv
from .graphconv import GraphConv
from .relgraphconv import RelGraphConv
from .sageconv import SAGEConv
from .sgconv import SGConv
python/dgl/nn/tensorflow/conv/appnpconv.py
+75
View File
@@ -0,0 +1,75 @@
"""TF Module for APPNPConv"""
# pylint: disable= no-member, arguments-differ, invalid-name
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from .... import function as fn
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.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.num_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
python/dgl/nn/tensorflow/conv/chebconv.py
+159
View File
@@ -0,0 +1,159 @@
"""Tensorflow Module for Chebyshev Spectral Graph Convolution layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from .... import broadcast_nodes, function as fn
from ....base import dgl_warning
class ChebConv(layers.Layer):
r"""Chebyshev Spectral Graph Convolution layer from `Convolutional
Neural Networks on Graphs with Fast Localized Spectral Filtering
<https://arxiv.org/pdf/1606.09375.pdf>`__
.. math::
h_i^{l+1} &= \sum_{k=0}^{K-1} W^{k, l}z_i^{k, l}
Z^{0, l} &= H^{l}
Z^{1, l} &= \tilde{L} \cdot H^{l}
Z^{k, l} &= 2 \cdot \tilde{L} \cdot Z^{k-1, l} - Z^{k-2, l}
\tilde{L} &= 2\left(I - \tilde{D}^{-1/2} \tilde{A} \tilde{D}^{-1/2}\right)/\lambda_{max} - I
where :math:`\tilde{A}` is :math:`A` + :math:`I`, :math:`W` is learnable weight.
Parameters
----------
in_feats: int
Dimension of input features; i.e, the number of dimensions of :math:`h_i^{(l)}`.
out_feats: int
Dimension of output features :math:`h_i^{(l+1)}`.
k : int
Chebyshev filter size :math:`K`.
activation : function, optional
Activation function. Default ``ReLu``.
bias : bool, optional
If True, adds a learnable bias to the output. Default: ``True``.
Example
-------
>>> import dgl
>>> import numpy as np
>>> import tensorflow as tf
>>> from dgl.nn import ChebConv
>>> with tf.device("CPU:0"):
... g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
... feat = tf.ones((6, 10))
... conv = ChebConv(10, 2, 2)
... res = conv(g, feat)
... res
<tf.Tensor: shape=(6, 2), dtype=float32, numpy=
array([[ 0.6163, -0.1809],
[ 0.6163, -0.1809],
[ 0.6163, -0.1809],
[ 0.9698, -1.5053],
[ 0.3664, 0.7556],
[-0.2370, 3.0164]], dtype=float32)>
"""
def __init__(
self, in_feats, out_feats, k, activation=tf.nn.relu, bias=True
):
super(ChebConv, self).__init__()
self._k = k
self._in_feats = in_feats
self._out_feats = out_feats
self.activation = activation
self.linear = layers.Dense(out_feats, use_bias=bias)
def call(self, graph, feat, lambda_max=None):
r"""Compute ChebNet layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor
The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
is size of input feature, :math:`N` is the number of nodes.
lambda_max : list or tensor or None, optional.
A list(tensor) with length :math:`B`, stores the largest eigenvalue
of the normalized laplacian of each individual graph in ``graph``,
where :math:`B` is the batch size of the input graph. Default: None.
If None, this method would set the default value to 2.
One can use :func:`dgl.laplacian_lambda_max` to compute this value.
Returns
-------
tf.Tensor
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
is size of output feature.
"""
def unnLaplacian(feat, D_invsqrt, graph):
"""Operation Feat * D^-1/2 A D^-1/2"""
graph.ndata["h"] = feat * D_invsqrt
graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
return graph.ndata.pop("h") * D_invsqrt
with graph.local_scope():
in_degrees = tf.clip_by_value(
tf.cast(graph.in_degrees(), tf.float32),
clip_value_min=1,
clip_value_max=np.inf,
)
D_invsqrt = tf.expand_dims(tf.pow(in_degrees, -0.5), axis=-1)
if lambda_max is None:
dgl_warning(
"lambda_max is not provided, using default value of 2. "
"Please use dgl.laplacian_lambda_max to compute the eigenvalues."
)
lambda_max = [2] * graph.batch_size
if isinstance(lambda_max, list):
lambda_max = tf.constant(lambda_max, dtype=tf.float32)
if lambda_max.ndim == 1:
lambda_max = tf.expand_dims(
lambda_max, axis=-1
) # (B,) to (B, 1)
# broadcast from (B, 1) to (N, 1)
lambda_max = broadcast_nodes(graph, lambda_max)
re_norm = 2.0 / lambda_max
# X_0 is the raw feature, Xt is the list of X_0, X_1, ... X_t
X_0 = feat
Xt = [X_0]
# X_1(f)
if self._k > 1:
h = unnLaplacian(X_0, D_invsqrt, graph)
X_1 = -re_norm * h + X_0 * (re_norm - 1)
# Append X_1 to Xt
Xt.append(X_1)
# Xi(x), i = 2...k
for _ in range(2, self._k):
h = unnLaplacian(X_1, D_invsqrt, graph)
X_i = -2 * re_norm * h + X_1 * 2 * (re_norm - 1) - X_0
# Append X_i to Xt
Xt.append(X_i)
X_1, X_0 = X_i, X_1
# Create the concatenation
Xt = tf.concat(Xt, 1)
# linear projection
h = self.linear(Xt)
# activation
if self.activation:
h = self.activation(h)
return h
python/dgl/nn/tensorflow/conv/densechebconv.py
+109
View File
@@ -0,0 +1,109 @@
"""Tensorflow Module for DenseChebConv"""
# pylint: disable= no-member, arguments-differ, invalid-name
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
class DenseChebConv(layers.Layer):
r"""Chebyshev Spectral Graph Convolution layer from `Convolutional
Neural Networks on Graphs with Fast Localized Spectral Filtering
<https://arxiv.org/pdf/1606.09375.pdf>`__
We recommend to use this module when applying ChebConv on dense graphs.
Parameters
----------
in_feats: int
Dimension of input features :math:`h_i^{(l)}`.
out_feats: int
Dimension of output features :math:`h_i^{(l+1)}`.
k : int
Chebyshev filter size.
activation : function, optional
Activation function, default is ReLu.
bias : bool, optional
If True, adds a learnable bias to the output. Default: ``True``.
See also
--------
`ChebConv <https://docs.dgl.ai/api/python/nn.tensorflow.html#chebconv>`__
"""
def __init__(self, in_feats, out_feats, k, bias=True):
super(DenseChebConv, self).__init__()
self._in_feats = in_feats
self._out_feats = out_feats
self._k = k
# keras initializer assume last two dims as fan_in and fan_out
xinit = tf.keras.initializers.glorot_normal()
self.W = tf.Variable(
initial_value=xinit(
shape=(k, in_feats, out_feats), dtype="float32"
),
trainable=True,
)
if bias:
zeroinit = tf.keras.initializers.zeros()
self.bias = tf.Variable(
initial_value=zeroinit(shape=(out_feats), dtype="float32"),
trainable=True,
)
else:
self.bias = None
def call(self, adj, feat, lambda_max=None):
r"""Compute (Dense) Chebyshev Spectral Graph Convolution layer.
Parameters
----------
adj : tf.Tensor
The adjacency matrix of the graph to apply Graph Convolution on,
should be of shape :math:`(N, N)`, where a row represents the destination
and a column represents the source.
feat : tf.Tensor
The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
is size of input feature, :math:`N` is the number of nodes.
lambda_max : float or None, optional
A float value indicates the largest eigenvalue of given graph.
Default: None.
Returns
-------
tf.Tensor
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
is size of output feature.
"""
A = adj
num_nodes = A.shape[0]
in_degree = 1 / tf.sqrt(
tf.clip_by_value(
tf.reduce_sum(A, 1), clip_value_min=1, clip_value_max=np.inf
)
)
D_invsqrt = tf.linalg.diag(in_degree)
I = tf.eye(num_nodes)
L = I - D_invsqrt @ A @ D_invsqrt
if lambda_max is None:
lambda_ = tf.linalg.eig(L)[0][:, 0]
lambda_max = tf.reduce_max(lambda_)
L_hat = 2 * L / lambda_max - I
Z = [tf.eye(num_nodes)]
for i in range(1, self._k):
if i == 1:
Z.append(L_hat)
else:
Z.append(2 * L_hat @ Z[-1] - Z[-2])
Zs = tf.stack(Z, 0) # (k, n, n)
Zh = Zs @ tf.expand_dims(feat, axis=0) @ self.W
Zh = tf.reduce_sum(Zh, 0)
if self.bias is not None:
Zh = Zh + self.bias
return Zh
python/dgl/nn/tensorflow/conv/edgeconv.py
+139
View File
@@ -0,0 +1,139 @@
"""Tensorflow modules for EdgeConv Layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import tensorflow as tf
from tensorflow.keras import layers
from .... import function as fn
from ....base import DGLError
from ....utils import expand_as_pair
class EdgeConv(layers.Layer):
r"""EdgeConv layer from `Dynamic Graph CNN for Learning on Point Clouds
<https://arxiv.org/pdf/1801.07829>`__
It can be described as follows:
.. math::
h_i^{(l+1)} = \max_{j \in \mathcal{N}(i)} (
\Theta \cdot (h_j^{(l)} - h_i^{(l)}) + \Phi \cdot h_i^{(l)})
where :math:`\mathcal{N}(i)` is the neighbor of :math:`i`,
:math:`\Theta` and :math:`\Phi` are linear layers.
.. note::
The original formulation includes a ReLU inside the maximum operator.
This is equivalent to first applying a maximum operator then applying
the ReLU.
Parameters
----------
in_feat : int
Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
out_feat : int
Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
batch_norm : bool
Whether to include batch normalization on messages. Default: ``False``.
allow_zero_in_degree : bool, optional
If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
since no message will be passed to those nodes. This is harmful for some applications
causing silent performance regression. This module will raise a DGLError if it detects
0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
and let the users handle it by themselves. Default: ``False``.
Note
----
Zero in-degree nodes will lead to invalid output value. This is because no message
will be passed to those nodes, the aggregation function will be appied on empty input.
A common practice to avoid this is to add a self-loop for each node in the graph if
it is homogeneous, which can be achieved by:
>>> g = ... # a DGLGraph
>>> g = dgl.add_self_loop(g)
Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
A common practise to handle this is to filter out the nodes with zero-in-degree when use
after conv.
"""
def __init__(self, out_feats, batch_norm=False, allow_zero_in_degree=False):
super(EdgeConv, self).__init__()
self.batch_norm = batch_norm
self._allow_zero_in_degree = allow_zero_in_degree
self.theta = layers.Dense(out_feats)
self.phi = layers.Dense(out_feats)
if batch_norm:
self.bn = layers.BatchNormalization()
def set_allow_zero_in_degree(self, set_value):
r"""Set allow_zero_in_degree flag.
Parameters
----------
set_value : bool
The value to be set to the flag.
"""
self._allow_zero_in_degree = set_value
def call(self, g, feat):
"""Forward computation
Parameters
----------
g : DGLGraph
The graph.
feat : tf.Tensor or pair of tf.Tensor
:math:`(N, D)` where :math:`N` is the number of nodes and
:math:`D` is the number of feature dimensions.
If a pair of tensors is given, the graph must be a uni-bipartite graph
with only one edge type, and the two tensors must have the same
dimensionality on all except the first axis.
Returns
-------
tf.Tensor or pair of tf.Tensor
New node features.
Raises
------
DGLError
If there are 0-in-degree nodes in the input graph, it will raise DGLError
since no message will be passed to those nodes. This will cause invalid output.
The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
"""
with g.local_scope():
if not self._allow_zero_in_degree:
if tf.math.count_nonzero(g.in_degrees() == 0) > 0:
raise DGLError(
"There are 0-in-degree nodes in the graph, "
"output for those nodes will be invalid. "
"This is harmful for some applications, "
"causing silent performance regression. "
"Adding self-loop on the input graph by "
"calling `g = dgl.add_self_loop(g)` will resolve "
"the issue. Setting ``allow_zero_in_degree`` "
"to be `True` when constructing this module will "
"suppress the check and let the code run."
)
h_src, h_dst = expand_as_pair(feat, g)
g.srcdata["x"] = h_src
g.dstdata["x"] = h_dst
g.apply_edges(fn.v_sub_u("x", "x", "theta"))
g.edata["theta"] = self.theta(g.edata["theta"])
g.dstdata["phi"] = self.phi(g.dstdata["x"])
if not self.batch_norm:
g.update_all(fn.e_add_v("theta", "phi", "e"), fn.max("e", "x"))
else:
g.apply_edges(fn.e_add_v("theta", "phi", "e"))
# for more comments on why global batch norm instead
# of batch norm within EdgeConv go to
# https://github.com/dmlc/dgl/blob/master/python/dgl/nn/pytorch/conv/edgeconv.py
g.edata["e"] = self.bn(g.edata["e"])
g.update_all(fn.copy_e("e", "e"), fn.max("e", "x"))
return g.dstdata["x"]
python/dgl/nn/tensorflow/conv/gatconv.py
+318
View File
@@ -0,0 +1,318 @@
"""Tensorflow modules for graph attention networks(GAT)."""
import numpy as np
# pylint: disable= no-member, arguments-differ, invalid-name
import tensorflow as tf
from tensorflow.keras import layers
from .... import function as fn
from ....base import DGLError
from ...functional import edge_softmax
from ..utils import Identity
# pylint: enable=W0235
class GATConv(layers.Layer):
r"""Graph Attention Layer from `Graph Attention Network
<https://arxiv.org/pdf/1710.10903.pdf>`__
.. math::
h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} \alpha_{i,j} W^{(l)} h_j^{(l)}
where :math:`\alpha_{ij}` is the attention score bewteen node :math:`i` and
node :math:`j`:
.. math::
\alpha_{ij}^{l} &= \mathrm{softmax_i} (e_{ij}^{l})
e_{ij}^{l} &= \mathrm{LeakyReLU}\left(\vec{a}^T [W h_{i} \| W h_{j}]\right)
Parameters
----------
in_feats : int, or pair of ints
Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
ATConv can be applied on homogeneous graph and unidirectional
`bipartite graph <https://docs.dgl.ai/generated/dgl.bipartite.html?highlight=bipartite>`__.
If the layer is to be applied to a unidirectional bipartite graph, ``in_feats``
specifies the input feature size on both the source and destination nodes. If
a scalar is given, the source and destination node feature size would take the
same value.
out_feats : int
Output feature size; i.e, the number of dimensions of :math:`h_i^{(l+1)}`.
num_heads : int
Number of heads in Multi-Head Attention.
feat_drop : float, optional
Dropout rate on feature. Defaults: ``0``.
attn_drop : float, optional
Dropout rate on attention weight. Defaults: ``0``.
negative_slope : float, optional
LeakyReLU angle of negative slope. Defaults: ``0.2``.
residual : bool, optional
If True, use residual connection. Defaults: ``False``.
activation : callable activation function/layer or None, optional.
If not None, applies an activation function to the updated node features.
Default: ``None``.
allow_zero_in_degree : bool, optional
If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
since no message will be passed to those nodes. This is harmful for some applications
causing silent performance regression. This module will raise a DGLError if it detects
0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
and let the users handle it by themselves. Defaults: ``False``.
Note
----
Zero in-degree nodes will lead to invalid output value. This is because no message
will be passed to those nodes, the aggregation function will be appied on empty input.
A common practice to avoid this is to add a self-loop for each node in the graph if
it is homogeneous, which can be achieved by:
>>> g = ... # a DGLGraph
>>> g = dgl.add_self_loop(g)
Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
A common practise to handle this is to filter out the nodes with zero-in-degree when use
after conv.
Examples
--------
>>> import dgl
>>> import numpy as np
>>> import tensorflow as tf
>>> from dgl.nn import GATConv
>>>
>>> # Case 1: Homogeneous graph
>>> with tf.device("CPU:0"):
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> g = dgl.add_self_loop(g)
>>> feat = tf.ones((6, 10))
>>> gatconv = GATConv(10, 2, num_heads=3)
>>> res = gatconv(g, feat)
>>> res
<tf.Tensor: shape=(6, 3, 2), dtype=float32, numpy=
array([[[ 0.75311995, -1.8093625 ],
[-0.12128812, -0.78072834],
[-0.49870574, -0.15074375]],
[[ 0.75311995, -1.8093625 ],
[-0.12128812, -0.78072834],
[-0.49870574, -0.15074375]],
[[ 0.75311995, -1.8093625 ],
[-0.12128812, -0.78072834],
[-0.49870574, -0.15074375]],
[[ 0.75311995, -1.8093626 ],
[-0.12128813, -0.78072834],
[-0.49870574, -0.15074375]],
[[ 0.75311995, -1.8093625 ],
[-0.12128812, -0.78072834],
[-0.49870574, -0.15074375]],
[[ 0.75311995, -1.8093625 ],
[-0.12128812, -0.78072834],
[-0.49870574, -0.15074375]]], dtype=float32)>
>>> # Case 2: Unidirectional bipartite graph
>>> u = [0, 1, 0, 0, 1]
>>> v = [0, 1, 2, 3, 2]
>>> g = dgl.heterograph({('A', 'r', 'B'): (u, v)})
>>> with tf.device("CPU:0"):
>>> u_feat = tf.convert_to_tensor(np.random.rand(2, 5))
>>> v_feat = tf.convert_to_tensor(np.random.rand(4, 10))
>>> gatconv = GATConv((5,10), 2, 3)
>>> res = gatconv(g, (u_feat, v_feat))
>>> res
<tf.Tensor: shape=(4, 3, 2), dtype=float32, numpy=
array([[[-0.89649093, -0.74841046],
[ 0.5088224 , 0.10908248],
[ 0.55670375, -0.6811229 ]],
[[-0.7905004 , -0.1457274 ],
[ 0.2248168 , 0.93014705],
[ 0.12816726, -0.4093595 ]],
[[-0.85875374, -0.53382933],
[ 0.36841977, 0.51498866],
[ 0.31893706, -0.5303393 ]],
[[-0.89649093, -0.74841046],
[ 0.5088224 , 0.10908248],
[ 0.55670375, -0.6811229 ]]], dtype=float32)>
"""
def __init__(
self,
in_feats,
out_feats,
num_heads,
feat_drop=0.0,
attn_drop=0.0,
negative_slope=0.2,
residual=False,
activation=None,
allow_zero_in_degree=False,
):
super(GATConv, self).__init__()
self._num_heads = num_heads
self._in_feats = in_feats
self._out_feats = out_feats
self._allow_zero_in_degree = allow_zero_in_degree
xinit = tf.keras.initializers.VarianceScaling(
scale=np.sqrt(2), mode="fan_avg", distribution="untruncated_normal"
)
if isinstance(in_feats, tuple):
self.fc_src = layers.Dense(
out_feats * num_heads, use_bias=False, kernel_initializer=xinit
)
self.fc_dst = layers.Dense(
out_feats * num_heads, use_bias=False, kernel_initializer=xinit
)
else:
self.fc = layers.Dense(
out_feats * num_heads, use_bias=False, kernel_initializer=xinit
)
self.attn_l = tf.Variable(
initial_value=xinit(
shape=(1, num_heads, out_feats), dtype="float32"
),
trainable=True,
)
self.attn_r = tf.Variable(
initial_value=xinit(
shape=(1, num_heads, out_feats), dtype="float32"
),
trainable=True,
)
self.feat_drop = layers.Dropout(rate=feat_drop)
self.attn_drop = layers.Dropout(rate=attn_drop)
self.leaky_relu = layers.LeakyReLU(alpha=negative_slope)
if residual:
if in_feats != out_feats:
self.res_fc = layers.Dense(
num_heads * out_feats,
use_bias=False,
kernel_initializer=xinit,
)
else:
self.res_fc = Identity()
else:
self.res_fc = None
# self.register_buffer('res_fc', None)
self.activation = activation
def set_allow_zero_in_degree(self, set_value):
r"""Set allow_zero_in_degree flag.
Parameters
----------
set_value : bool
The value to be set to the flag.
"""
self._allow_zero_in_degree = set_value
def call(self, graph, feat, get_attention=False):
r"""Compute graph attention network layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor or pair of tf.Tensor
If a tf.Tensor is given, the input feature of shape :math:`(N, *, D_{in})` where
:math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
If a pair of tf.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, *, D_{in_{src}})` and :math:`(N_{out}, *, D_{in_{dst}})`.
get_attention : bool, optional
Whether to return the attention values. Default to False.
Returns
-------
tf.Tensor
The output feature of shape :math:`(N, *, H, D_{out})` where :math:`H`
is the number of heads, and :math:`D_{out}` is size of output feature.
tf.Tensor, optional
The attention values of shape :math:`(E, *, H, 1)`, where :math:`E` is the number of
edges. This is returned only when :attr:`get_attention` is ``True``.
Raises
------
DGLError
If there are 0-in-degree nodes in the input graph, it will raise DGLError
since no message will be passed to those nodes. This will cause invalid output.
The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
"""
with graph.local_scope():
if not self._allow_zero_in_degree:
if tf.math.count_nonzero(graph.in_degrees() == 0) > 0:
raise DGLError(
"There are 0-in-degree nodes in the graph, "
"output for those nodes will be invalid. "
"This is harmful for some applications, "
"causing silent performance regression. "
"Adding self-loop on the input graph by "
"calling `g = dgl.add_self_loop(g)` will resolve "
"the issue. Setting ``allow_zero_in_degree`` "
"to be `True` when constructing this module will "
"suppress the check and let the code run."
)
if isinstance(feat, tuple):
src_prefix_shape = tuple(feat[0].shape[:-1])
dst_prefix_shape = tuple(feat[1].shape[:-1])
h_src = self.feat_drop(feat[0])
h_dst = self.feat_drop(feat[1])
if not hasattr(self, "fc_src"):
self.fc_src, self.fc_dst = self.fc, self.fc
feat_src = tf.reshape(
self.fc_src(h_src),
src_prefix_shape + (self._num_heads, self._out_feats),
)
feat_dst = tf.reshape(
self.fc_dst(h_dst),
dst_prefix_shape + (self._num_heads, self._out_feats),
)
else:
src_prefix_shape = dst_prefix_shape = tuple(feat.shape[:-1])
h_src = h_dst = self.feat_drop(feat)
feat_src = feat_dst = tf.reshape(
self.fc(h_src),
src_prefix_shape + (self._num_heads, self._out_feats),
)
if graph.is_block:
feat_dst = feat_src[: graph.number_of_dst_nodes()]
h_dst = h_dst[: graph.number_of_dst_nodes()]
dst_prefix_shape = (
graph.number_of_dst_nodes(),
) + dst_prefix_shape[1:]
# NOTE: GAT paper uses "first concatenation then linear projection"
# to compute attention scores, while ours is "first projection then
# addition", the two approaches are mathematically equivalent:
# We decompose the weight vector a mentioned in the paper into
# [a_l || a_r], then
# a^T [Wh_i || Wh_j] = a_l Wh_i + a_r Wh_j
# Our implementation is much efficient because we do not need to
# save [Wh_i || Wh_j] on edges, which is not memory-efficient. Plus,
# addition could be optimized with DGL's built-in function u_add_v,
# which further speeds up computation and saves memory footprint.
el = tf.reduce_sum(feat_src * self.attn_l, axis=-1, keepdims=True)
er = tf.reduce_sum(feat_dst * self.attn_r, axis=-1, keepdims=True)
graph.srcdata.update({"ft": feat_src, "el": el})
graph.dstdata.update({"er": er})
# compute edge attention, el and er are a_l Wh_i and a_r Wh_j respectively.
graph.apply_edges(fn.u_add_v("el", "er", "e"))
e = self.leaky_relu(graph.edata.pop("e"))
# compute softmax
graph.edata["a"] = self.attn_drop(edge_softmax(graph, e))
# message passing
graph.update_all(fn.u_mul_e("ft", "a", "m"), fn.sum("m", "ft"))
rst = graph.dstdata["ft"]
# residual
if self.res_fc is not None:
resval = tf.reshape(
self.res_fc(h_dst), dst_prefix_shape + (-1, self._out_feats)
)
rst = rst + resval
# activation
if self.activation:
rst = self.activation(rst)
if get_attention:
return rst, graph.edata["a"]
else:
return rst
python/dgl/nn/tensorflow/conv/ginconv.py
+111
View File
@@ -0,0 +1,111 @@
"""Tensorflow Module for Graph Isomorphism Network layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import tensorflow as tf
from tensorflow.keras import layers
from .... import function as fn
from ....utils import expand_as_pair
class GINConv(layers.Layer):
r"""Graph Isomorphism Network layer from `How Powerful are Graph
Neural Networks? <https://arxiv.org/pdf/1810.00826.pdf>`__
.. math::
h_i^{(l+1)} = f_\Theta \left((1 + \epsilon) h_i^{l} +
\mathrm{aggregate}\left(\left\{h_j^{l}, j\in\mathcal{N}(i)
\right\}\right)\right)
Parameters
----------
apply_func : callable activation function/layer or None
If not None, apply this function to the updated node feature,
the :math:`f_\Theta` in the formula.
aggregator_type : str
Aggregator type to use (``sum``, ``max`` or ``mean``).
init_eps : float, optional
Initial :math:`\epsilon` value, default: ``0``.
learn_eps : bool, optional
If True, :math:`\epsilon` will be a learnable parameter. Default: ``False``.
Example
-------
>>> import dgl
>>> import numpy as np
>>> import tensorflow as tf
>>> from dgl.nn import GINConv
>>>
>>> with tf.device("CPU:0"):
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = tf.ones((6, 10))
>>> lin = tf.keras.layers.Dense(10)
>>> conv = GINConv(lin, 'max')
>>> res = conv(g, feat)
>>> res
<tf.Tensor: shape=(6, 10), dtype=float32, numpy=
array([[-0.1090256 , 1.9050574 , -0.30704725, -1.995831 , -0.36399186,
1.10414 , 2.4885745 , -0.35387516, 1.3568261 , 1.7267858 ],
[-0.1090256 , 1.9050574 , -0.30704725, -1.995831 , -0.36399186,
1.10414 , 2.4885745 , -0.35387516, 1.3568261 , 1.7267858 ],
[-0.1090256 , 1.9050574 , -0.30704725, -1.995831 , -0.36399186,
1.10414 , 2.4885745 , -0.35387516, 1.3568261 , 1.7267858 ],
[-0.1090256 , 1.9050574 , -0.30704725, -1.995831 , -0.36399186,
1.10414 , 2.4885745 , -0.35387516, 1.3568261 , 1.7267858 ],
[-0.1090256 , 1.9050574 , -0.30704725, -1.995831 , -0.36399186,
1.10414 , 2.4885745 , -0.35387516, 1.3568261 , 1.7267858 ],
[-0.0545128 , 0.9525287 , -0.15352362, -0.9979155 , -0.18199593,
0.55207 , 1.2442873 , -0.17693758, 0.67841303, 0.8633929 ]],
dtype=float32)>
"""
def __init__(
self, apply_func, aggregator_type, init_eps=0, learn_eps=False
):
super(GINConv, self).__init__()
self.apply_func = apply_func
if aggregator_type == "sum":
self._reducer = fn.sum
elif aggregator_type == "max":
self._reducer = fn.max
elif aggregator_type == "mean":
self._reducer = fn.mean
else:
raise KeyError(
"Aggregator type {} not recognized.".format(aggregator_type)
)
# to specify whether eps is trainable or not.
self.eps = tf.Variable(
initial_value=[init_eps], dtype=tf.float32, trainable=learn_eps
)
def call(self, graph, feat):
r"""Compute Graph Isomorphism Network layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor or pair of tf.Tensor
If a tf.Tensor is given, the input feature of shape :math:`(N, D_{in})` where
:math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
If a pair of tf.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, D_{in})` and :math:`(N_{out}, D_{in})`.
If ``apply_func`` is not None, :math:`D_{in}` should
fit the input dimensionality requirement of ``apply_func``.
Returns
-------
tf.Tensor
The output feature of shape :math:`(N, D_{out})` where
:math:`D_{out}` is the output dimensionality of ``apply_func``.
If ``apply_func`` is None, :math:`D_{out}` should be the same
as input dimensionality.
"""
with graph.local_scope():
feat_src, feat_dst = expand_as_pair(feat, graph)
graph.srcdata["h"] = feat_src
graph.update_all(fn.copy_u("h", "m"), self._reducer("m", "neigh"))
rst = (1 + self.eps) * feat_dst + graph.dstdata["neigh"]
if self.apply_func is not None:
rst = self.apply_func(rst)
return rst
python/dgl/nn/tensorflow/conv/graphconv.py
+318
View File
@@ -0,0 +1,318 @@
"""Tensorflow modules for graph convolutions(GCN)."""
import numpy as np
# pylint: disable= no-member, arguments-differ, invalid-name
import tensorflow as tf
from tensorflow.keras import layers
from .... import function as fn
from ....base import DGLError
from ....utils import expand_as_pair
# pylint: disable=W0235
class GraphConv(layers.Layer):
r"""Graph convolution from `Semi-Supervised Classification with Graph Convolutional Networks
<https://arxiv.org/abs/1609.02907>`__
Mathematically it is defined as follows:
.. math::
h_i^{(l+1)} = \sigma(b^{(l)} + \sum_{j\in\mathcal{N}(i)}\frac{1}{c_{ij}}h_j^{(l)}W^{(l)})
where :math:`\mathcal{N}(i)` is the set of neighbors of node :math:`i`,
:math:`c_{ij}` is the product of the square root of node degrees
(i.e., :math:`c_{ij} = \sqrt{|\mathcal{N}(i)|}\sqrt{|\mathcal{N}(j)|}`),
and :math:`\sigma` is an activation function.
Parameters
----------
in_feats : int
Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
out_feats : int
Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
norm : str, optional
How to apply the normalizer. Can be one of the following values:
* ``right``, to divide the aggregated messages by each node's in-degrees,
which is equivalent to averaging the received messages.
* ``none``, where no normalization is applied.
* ``both`` (default), where the messages are scaled with :math:`1/c_{ji}` above, equivalent
to symmetric normalization.
* ``left``, to divide the messages sent out from each node by its out-degrees,
equivalent to random walk normalization.
weight : bool, optional
If True, apply a linear layer. Otherwise, aggregating the messages
without a weight matrix.
bias : bool, optional
If True, adds a learnable bias to the output. Default: ``True``.
activation : callable activation function/layer or None, optional
If not None, applies an activation function to the updated node features.
Default: ``None``.
allow_zero_in_degree : bool, optional
If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
since no message will be passed to those nodes. This is harmful for some applications
causing silent performance regression. This module will raise a DGLError if it detects
0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
and let the users handle it by themselves. Default: ``False``.
Attributes
----------
weight : torch.Tensor
The learnable weight tensor.
bias : torch.Tensor
The learnable bias tensor.
Note
----
Zero in-degree nodes will lead to invalid output value. This is because no message
will be passed to those nodes, the aggregation function will be appied on empty input.
A common practice to avoid this is to add a self-loop for each node in the graph if
it is homogeneous, which can be achieved by:
>>> g = ... # a DGLGraph
>>> g = dgl.add_self_loop(g)
Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
A common practise to handle this is to filter out the nodes with zero-in-degree when use
after conv.
Examples
--------
>>> import dgl
>>> import numpy as np
>>> import tensorflow as tf
>>> from dgl.nn import GraphConv
>>> # Case 1: Homogeneous graph
>>> with tf.device("CPU:0"):
... g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
... g = dgl.add_self_loop(g)
... feat = tf.ones((6, 10))
... conv = GraphConv(10, 2, norm='both', weight=True, bias=True)
... res = conv(g, feat)
>>> print(res)
<tf.Tensor: shape=(6, 2), dtype=float32, numpy=
array([[ 0.6208475 , -0.4896223 ],
[ 0.68356586, -0.5390842 ],
[ 0.6208475 , -0.4896223 ],
[ 0.7859846 , -0.61985517],
[ 0.8251371 , -0.65073216],
[ 0.48335412, -0.38119012]], dtype=float32)>
>>> # allow_zero_in_degree example
>>> with tf.device("CPU:0"):
... g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
... conv = GraphConv(10, 2, norm='both', weight=True, bias=True, allow_zero_in_degree=True)
... res = conv(g, feat)
>>> print(res)
<tf.Tensor: shape=(6, 2), dtype=float32, numpy=
array([[ 0.6208475 , -0.4896223 ],
[ 0.68356586, -0.5390842 ],
[ 0.6208475 , -0.4896223 ],
[ 0.7859846 , -0.61985517],
[ 0.8251371 , -0.65073216],
[ 0., 0.]], dtype=float32)>
>>> # Case 2: Unidirectional bipartite graph
>>> u = [0, 1, 0, 0, 1]
>>> v = [0, 1, 2, 3, 2]
>>> with tf.device("CPU:0"):
... g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
... u_fea = tf.convert_to_tensor(np.random.rand(2, 5))
... v_fea = tf.convert_to_tensor(np.random.rand(4, 5))
... conv = GraphConv(5, 2, norm='both', weight=True, bias=True)
... res = conv(g, (u_fea, v_fea))
>>> res
<tf.Tensor: shape=(4, 2), dtype=float32, numpy=
array([[ 1.3607183, -0.1636453],
[ 1.6665325, -0.2004239],
[ 2.1405895, -0.2574358],
[ 1.3607183, -0.1636453]], dtype=float32)>
"""
def __init__(
self,
in_feats,
out_feats,
norm="both",
weight=True,
bias=True,
activation=None,
allow_zero_in_degree=False,
):
super(GraphConv, self).__init__()
if norm not in ("none", "both", "right", "left"):
raise DGLError(
'Invalid norm value. Must be either "none", "both", "right" or "left".'
' But got "{}".'.format(norm)
)
self._in_feats = in_feats
self._out_feats = out_feats
self._norm = norm
self._allow_zero_in_degree = allow_zero_in_degree
if weight:
xinit = tf.keras.initializers.glorot_uniform()
self.weight = tf.Variable(
initial_value=xinit(
shape=(in_feats, out_feats), dtype="float32"
),
trainable=True,
)
else:
self.weight = None
if bias:
zeroinit = tf.keras.initializers.zeros()
self.bias = tf.Variable(
initial_value=zeroinit(shape=(out_feats), dtype="float32"),
trainable=True,
)
else:
self.bias = None
self._activation = activation
def set_allow_zero_in_degree(self, set_value):
r"""Set allow_zero_in_degree flag.
Parameters
----------
set_value : bool
The value to be set to the flag.
"""
self._allow_zero_in_degree = set_value
def call(self, graph, feat, weight=None):
r"""Compute graph convolution.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor or pair of torch.Tensor
If a torch.Tensor is given, it represents the input feature of shape
:math:`(N, D_{in})`
where :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
If a pair of torch.Tensor is given, which is the case for bipartite graph, the pair
must contain two tensors of shape :math:`(N_{in}, D_{in_{src}})` and
:math:`(N_{out}, D_{in_{dst}})`.
weight : torch.Tensor, optional
Optional external weight tensor.
Returns
-------
torch.Tensor
The output feature
Raises
------
DGLError
If there are 0-in-degree nodes in the input graph, it will raise DGLError
since no message will be passed to those nodes. This will cause invalid output.
The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
Note
----
* Input shape: :math:`(N, *, \text{in_feats})` where * means any number of additional
dimensions, :math:`N` is the number of nodes.
* Output shape: :math:`(N, *, \text{out_feats})` where all but the last dimension are
the same shape as the input.
* Weight shape: :math:`(\text{in_feats}, \text{out_feats})`.
"""
with graph.local_scope():
if not self._allow_zero_in_degree:
if tf.math.count_nonzero(graph.in_degrees() == 0) > 0:
raise DGLError(
"There are 0-in-degree nodes in the graph, "
"output for those nodes will be invalid. "
"This is harmful for some applications, "
"causing silent performance regression. "
"Adding self-loop on the input graph by "
"calling `g = dgl.add_self_loop(g)` will resolve "
"the issue. Setting ``allow_zero_in_degree`` "
"to be `True` when constructing this module will "
"suppress the check and let the code run."
)
feat_src, feat_dst = expand_as_pair(feat, graph)
if self._norm in ["both", "left"]:
degs = tf.clip_by_value(
tf.cast(graph.out_degrees(), tf.float32),
clip_value_min=1,
clip_value_max=np.inf,
)
if self._norm == "both":
norm = tf.pow(degs, -0.5)
else:
norm = 1.0 / degs
shp = norm.shape + (1,) * (feat_dst.ndim - 1)
norm = tf.reshape(norm, shp)
feat_src = feat_src * norm
if weight is not None:
if self.weight is not None:
raise DGLError(
"External weight is provided while at the same time the"
" module has defined its own weight parameter. Please"
" create the module with flag weight=False."
)
else:
weight = self.weight
if self._in_feats > self._out_feats:
# mult W first to reduce the feature size for aggregation.
if weight is not None:
feat_src = tf.matmul(feat_src, weight)
graph.srcdata["h"] = feat_src
graph.update_all(
fn.copy_u(u="h", out="m"), fn.sum(msg="m", out="h")
)
rst = graph.dstdata["h"]
else:
# aggregate first then mult W
graph.srcdata["h"] = feat_src
graph.update_all(
fn.copy_u(u="h", out="m"), fn.sum(msg="m", out="h")
)
rst = graph.dstdata["h"]
if weight is not None:
rst = tf.matmul(rst, weight)
if self._norm in ["both", "right"]:
degs = tf.clip_by_value(
tf.cast(graph.in_degrees(), tf.float32),
clip_value_min=1,
clip_value_max=np.inf,
)
if self._norm == "both":
norm = tf.pow(degs, -0.5)
else:
norm = 1.0 / degs
shp = norm.shape + (1,) * (feat_dst.ndim - 1)
norm = tf.reshape(norm, shp)
rst = rst * norm
if self.bias is not None:
rst = rst + self.bias
if self._activation is not None:
rst = self._activation(rst)
return rst
def extra_repr(self):
"""Set the extra representation of the module,
which will come into effect when printing the model.
"""
summary = "in={_in_feats}, out={_out_feats}"
summary += ", normalization={_norm}"
if "_activation" in self.__dict__:
summary += ", activation={_activation}"
return summary.format(**self.__dict__)
python/dgl/nn/tensorflow/conv/relgraphconv.py
+330
View File
@@ -0,0 +1,330 @@
"""Tensorflow Module for Relational graph convolution layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import tensorflow as tf
from tensorflow.keras import layers
from .... import function as fn
from .. import utils
class RelGraphConv(layers.Layer):
r"""Relational graph convolution layer from `Modeling Relational Data with Graph
Convolutional Networks <https://arxiv.org/abs/1703.06103>`__
It can be described as below:
.. math::
h_i^{(l+1)} = \sigma(\sum_{r\in\mathcal{R}}
\sum_{j\in\mathcal{N}^r(i)}\frac{1}{c_{i,r}}W_r^{(l)}h_j^{(l)}+W_0^{(l)}h_i^{(l)})
where :math:`\mathcal{N}^r(i)` is the neighbor set of node :math:`i` w.r.t. relation
:math:`r`. :math:`c_{i,r}` is the normalizer equal
to :math:`|\mathcal{N}^r(i)|`. :math:`\sigma` is an activation function. :math:`W_0`
is the self-loop weight.
The basis regularization decomposes :math:`W_r` by:
.. math::
W_r^{(l)} = \sum_{b=1}^B a_{rb}^{(l)}V_b^{(l)}
where :math:`B` is the number of bases, :math:`V_b^{(l)}` are linearly combined
with coefficients :math:`a_{rb}^{(l)}`.
The block-diagonal-decomposition regularization decomposes :math:`W_r` into :math:`B`
number of block diagonal matrices. We refer :math:`B` as the number of bases.
The block regularization decomposes :math:`W_r` by:
.. math::
W_r^{(l)} = \oplus_{b=1}^B Q_{rb}^{(l)}
where :math:`B` is the number of bases, :math:`Q_{rb}^{(l)}` are block
bases with shape :math:`R^{(d^{(l+1)}/B)*(d^{l}/B)}`.
Parameters
----------
in_feat : int
Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
out_feat : int
Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
num_rels : int
Number of relations. .
regularizer : str
Which weight regularizer to use "basis" or "bdd".
"basis" is short for basis-diagonal-decomposition.
"bdd" is short for block-diagonal-decomposition.
num_bases : int, optional
Number of bases. If is none, use number of relations. Default: ``None``.
bias : bool, optional
True if bias is added. Default: ``True``.
activation : callable, optional
Activation function. Default: ``None``.
self_loop : bool, optional
True to include self loop message. Default: ``True``.
low_mem : bool, optional
True to use low memory implementation of relation message passing function. Default: False.
This option trades speed with memory consumption, and will slowdown the forward/backward.
Turn it on when you encounter OOM problem during training or evaluation. Default: ``False``.
dropout : float, optional
Dropout rate. Default: ``0.0``
layer_norm: float, optional
Add layer norm. Default: ``False``
Examples
--------
>>> import dgl
>>> import numpy as np
>>> import tensorflow as tf
>>> from dgl.nn import RelGraphConv
>>>
>>> with tf.device("CPU:0"):
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = tf.ones((6, 10))
>>> conv = RelGraphConv(10, 2, 3, regularizer='basis', num_bases=2)
>>> etype = tf.convert_to_tensor(np.array([0,1,2,0,1,2]).astype(np.int64))
>>> res = conv(g, feat, etype)
>>> res
<tf.Tensor: shape=(6, 2), dtype=float32, numpy=
array([[-0.02938664, 1.7932655 ],
[ 0.1146394 , 0.48319 ],
[-0.02938664, 1.7932655 ],
[ 1.2054908 , -0.26098895],
[ 0.1146394 , 0.48319 ],
[ 0.75915515, 1.1454091 ]], dtype=float32)>
>>> # One-hot input
>>> with tf.device("CPU:0"):
>>> one_hot_feat = tf.convert_to_tensor(np.array([0,1,2,3,4,5]).astype(np.int64))
>>> res = conv(g, one_hot_feat, etype)
>>> res
<tf.Tensor: shape=(6, 2), dtype=float32, numpy=
array([[-0.24205256, -0.7922753 ],
[ 0.62085056, 0.4893622 ],
[-0.9484881 , -0.26546806],
[-0.2163915 , -0.12585883],
[-0.14293689, 0.77483284],
[ 0.091169 , -0.06761569]], dtype=float32)>
"""
def __init__(
self,
in_feat,
out_feat,
num_rels,
regularizer="basis",
num_bases=None,
bias=True,
activation=None,
self_loop=True,
low_mem=False,
dropout=0.0,
layer_norm=False,
):
super(RelGraphConv, self).__init__()
self.in_feat = in_feat
self.out_feat = out_feat
self.num_rels = num_rels
self.regularizer = regularizer
self.num_bases = num_bases
if (
self.num_bases is None
or self.num_bases > self.num_rels
or self.num_bases < 0
):
self.num_bases = self.num_rels
self.bias = bias
self.activation = activation
self.self_loop = self_loop
self.low_mem = low_mem
assert (
layer_norm is False
), "TensorFlow currently does not support layer norm."
xinit = tf.keras.initializers.glorot_uniform()
zeroinit = tf.keras.initializers.zeros()
if regularizer == "basis":
# add basis weights
self.weight = tf.Variable(
initial_value=xinit(
shape=(self.num_bases, self.in_feat, self.out_feat),
dtype="float32",
),
trainable=True,
)
if self.num_bases < self.num_rels:
# linear combination coefficients
self.w_comp = tf.Variable(
initial_value=xinit(
shape=(self.num_rels, self.num_bases), dtype="float32"
),
trainable=True,
)
# message func
self.message_func = self.basis_message_func
elif regularizer == "bdd":
if in_feat % num_bases != 0 or out_feat % num_bases != 0:
raise ValueError(
"Feature size must be a multiplier of num_bases."
)
# add block diagonal weights
self.submat_in = in_feat // self.num_bases
self.submat_out = out_feat // self.num_bases
# assuming in_feat and out_feat are both divisible by num_bases
self.weight = tf.Variable(
initial_value=xinit(
shape=(
self.num_rels,
self.num_bases * self.submat_in * self.submat_out,
),
dtype="float32",
),
trainable=True,
)
# message func
self.message_func = self.bdd_message_func
else:
raise ValueError("Regularizer must be either 'basis' or 'bdd'")
# bias
if self.bias:
self.h_bias = tf.Variable(
initial_value=zeroinit(shape=(out_feat), dtype="float32"),
trainable=True,
)
# weight for self loop
if self.self_loop:
self.loop_weight = tf.Variable(
initial_value=xinit(shape=(in_feat, out_feat), dtype="float32"),
trainable=True,
)
self.dropout = layers.Dropout(rate=dropout)
def basis_message_func(self, edges):
"""Message function for basis regularizer"""
if self.num_bases < self.num_rels:
# generate all weights from bases
weight = tf.reshape(
self.weight, (self.num_bases, self.in_feat * self.out_feat)
)
weight = tf.reshape(
tf.matmul(self.w_comp, weight),
(self.num_rels, self.in_feat, self.out_feat),
)
else:
weight = self.weight
# calculate msg @ W_r before put msg into edge
# if src is th.int64 we expect it is an index select
if edges.src["h"].dtype != tf.int64 and self.low_mem:
etypes, _ = tf.unique(edges.data["type"])
msg = tf.zeros([edges.src["h"].shape[0], self.out_feat])
idx = tf.range(edges.src["h"].shape[0])
for etype in etypes:
loc = edges.data["type"] == etype
w = weight[etype]
src = tf.boolean_mask(edges.src["h"], loc)
sub_msg = tf.matmul(src, w)
indices = tf.reshape(tf.boolean_mask(idx, loc), (-1, 1))
msg = tf.tensor_scatter_nd_update(msg, indices, sub_msg)
else:
msg = utils.bmm_maybe_select(
edges.src["h"], weight, edges.data["type"]
)
if "norm" in edges.data:
msg = msg * edges.data["norm"]
return {"msg": msg}
def bdd_message_func(self, edges):
"""Message function for block-diagonal-decomposition regularizer"""
if (edges.src["h"].dtype == tf.int64) and len(
edges.src["h"].shape
) == 1:
raise TypeError(
"Block decomposition does not allow integer ID feature."
)
# calculate msg @ W_r before put msg into edge
# if src is th.int64 we expect it is an index select
if self.low_mem:
etypes, _ = tf.unique(edges.data["type"])
msg = tf.zeros([edges.src["h"].shape[0], self.out_feat])
idx = tf.range(edges.src["h"].shape[0])
for etype in etypes:
loc = edges.data["type"] == etype
w = tf.reshape(
self.weight[etype],
(self.num_bases, self.submat_in, self.submat_out),
)
src = tf.reshape(
tf.boolean_mask(edges.src["h"], loc),
(-1, self.num_bases, self.submat_in),
)
sub_msg = tf.einsum("abc,bcd->abd", src, w)
sub_msg = tf.reshape(sub_msg, (-1, self.out_feat))
indices = tf.reshape(tf.boolean_mask(idx, loc), (-1, 1))
msg = tf.tensor_scatter_nd_update(msg, indices, sub_msg)
else:
weight = tf.reshape(
tf.gather(self.weight, edges.data["type"]),
(-1, self.submat_in, self.submat_out),
)
node = tf.reshape(edges.src["h"], (-1, 1, self.submat_in))
msg = tf.reshape(tf.matmul(node, weight), (-1, self.out_feat))
if "norm" in edges.data:
msg = msg * edges.data["norm"]
return {"msg": msg}
def call(self, g, x, etypes, norm=None):
"""Forward computation
Parameters
----------
g : DGLGraph
The graph.
x : tf.Tensor
Input node features. Could be either
* :math:`(|V|, D)` dense tensor
* :math:`(|V|,)` int64 vector, representing the categorical values of each
node. We then treat the input feature as an one-hot encoding feature.
etypes : tf.Tensor
Edge type tensor. Shape: :math:`(|E|,)`
norm : tf.Tensor
Optional edge normalizer tensor. Shape: :math:`(|E|, 1)`
Returns
-------
tf.Tensor
New node features.
"""
assert g.is_homogeneous, (
"not a homogeneous graph; convert it with to_homogeneous "
"and pass in the edge type as argument"
)
with g.local_scope():
g.ndata["h"] = x
g.edata["type"] = tf.cast(etypes, tf.int64)
if norm is not None:
g.edata["norm"] = norm
if self.self_loop:
loop_message = utils.matmul_maybe_select(x, self.loop_weight)
# message passing
g.update_all(self.message_func, fn.sum(msg="msg", out="h"))
# apply bias and activation
node_repr = g.ndata["h"]
if self.bias:
node_repr = node_repr + self.h_bias
if self.self_loop:
node_repr = node_repr + loop_message
if self.activation:
node_repr = self.activation(node_repr)
node_repr = self.dropout(node_repr)
return node_repr
python/dgl/nn/tensorflow/conv/sageconv.py
+213
View File
@@ -0,0 +1,213 @@
"""Tensorflow Module for GraphSAGE layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import tensorflow as tf
from tensorflow.keras import layers
from .... import function as fn
from ....base import DGLError
from ....utils import check_eq_shape, expand_as_pair
class SAGEConv(layers.Layer):
r"""GraphSAGE layer from `Inductive Representation Learning on
Large Graphs <https://arxiv.org/pdf/1706.02216.pdf>`__
.. math::
h_{\mathcal{N}(i)}^{(l+1)} &= \mathrm{aggregate}
\left(\{h_{j}^{l}, \forall j \in \mathcal{N}(i) \}\right)
h_{i}^{(l+1)} &= \sigma \left(W \cdot \mathrm{concat}
(h_{i}^{l}, h_{\mathcal{N}(i)}^{l+1}) \right)
h_{i}^{(l+1)} &= \mathrm{norm}(h_{i}^{(l+1)})
Parameters
----------
in_feats : int, or pair of ints
Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
GATConv can be applied on homogeneous graph and unidirectional
`bipartite graph <https://docs.dgl.ai/generated/dgl.bipartite.html?highlight=bipartite>`__.
If the layer applies on a unidirectional bipartite graph, ``in_feats``
specifies the input feature size on both the source and destination nodes. If
a scalar is given, the source and destination node feature size would take the
same value.
If aggregator type is ``gcn``, the feature size of source and destination nodes
are required to be the same.
out_feats : int
Output feature size; i.e, the number of dimensions of :math:`h_i^{(l+1)}`.
aggregator_type : str
Aggregator type to use (``mean``, ``gcn``, ``pool``, ``lstm``).
feat_drop : float
Dropout rate on features, default: ``0``.
bias : bool
If True, adds a learnable bias to the output. Default: ``True``.
norm : callable activation function/layer or None, optional
If not None, applies normalization to the updated node features.
activation : callable activation function/layer or None, optional
If not None, applies an activation function to the updated node features.
Default: ``None``.
Examples
--------
>>> import dgl
>>> import numpy as np
>>> import tensorflow as tf
>>> from dgl.nn import SAGEConv
>>>
>>> # Case 1: Homogeneous graph
>>> with tf.device("CPU:0"):
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> g = dgl.add_self_loop(g)
>>> feat = tf.ones((6, 10))
>>> conv = SAGEConv(10, 2, 'pool')
>>> res = conv(g, feat)
>>> res
<tf.Tensor: shape=(6, 2), dtype=float32, numpy=
array([[-3.6633523 , -0.90711546],
[-3.6633523 , -0.90711546],
[-3.6633523 , -0.90711546],
[-3.6633523 , -0.90711546],
[-3.6633523 , -0.90711546],
[-3.6633523 , -0.90711546]], dtype=float32)>
>>> # Case 2: Unidirectional bipartite graph
>>> with tf.device("CPU:0"):
>>> u = [0, 1, 0, 0, 1]
>>> v = [0, 1, 2, 3, 2]
>>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
>>> u_fea = tf.convert_to_tensor(np.random.rand(2, 5))
>>> v_fea = tf.convert_to_tensor(np.random.rand(4, 5))
>>> conv = SAGEConv((5, 10), 2, 'mean')
>>> res = conv(g, (u_fea, v_fea))
>>> res
<tf.Tensor: shape=(4, 2), dtype=float32, numpy=
array([[-0.59453356, -0.4055441 ],
[-0.47459763, -0.717764 ],
[ 0.3221837 , -0.29876417],
[-0.63356155, 0.09390211]], dtype=float32)>
"""
def __init__(
self,
in_feats,
out_feats,
aggregator_type,
feat_drop=0.0,
bias=True,
norm=None,
activation=None,
):
super(SAGEConv, self).__init__()
valid_aggre_types = {"mean", "gcn", "pool", "lstm"}
if aggregator_type not in valid_aggre_types:
raise DGLError(
"Invalid aggregator_type. Must be one of {}. "
"But got {!r} instead.".format(
valid_aggre_types, aggregator_type
)
)
self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
self._out_feats = out_feats
self._aggre_type = aggregator_type
self.norm = norm
self.feat_drop = layers.Dropout(feat_drop)
self.activation = activation
# aggregator type: mean/pool/lstm/gcn
if aggregator_type == "pool":
self.fc_pool = layers.Dense(self._in_src_feats)
if aggregator_type == "lstm":
self.lstm = layers.LSTM(units=self._in_src_feats)
if aggregator_type != "gcn":
self.fc_self = layers.Dense(out_feats, use_bias=bias)
self.fc_neigh = layers.Dense(out_feats, use_bias=bias)
def _lstm_reducer(self, nodes):
"""LSTM reducer
NOTE(zihao): lstm reducer with default schedule (degree bucketing)
is slow, we could accelerate this with degree padding in the future.
"""
m = nodes.mailbox["m"] # (B, L, D)
rst = self.lstm(m)
return {"neigh": rst}
def call(self, graph, feat):
r"""Compute GraphSAGE layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor or pair of tf.Tensor
If a tf.Tensor is given, it represents the input feature of shape
:math:`(N, D_{in})`
where :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
If a pair of tf.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
Returns
-------
tf.Tensor
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
is size of output feature.
"""
with graph.local_scope():
if isinstance(feat, tuple):
feat_src = self.feat_drop(feat[0])
feat_dst = self.feat_drop(feat[1])
else:
feat_src = feat_dst = self.feat_drop(feat)
if graph.is_block:
feat_dst = feat_src[: graph.number_of_dst_nodes()]
h_self = feat_dst
# Handle the case of graphs without edges
if graph.num_edges() == 0:
graph.dstdata["neigh"] = tf.cast(
tf.zeros((graph.number_of_dst_nodes(), self._in_src_feats)),
tf.float32,
)
if self._aggre_type == "mean":
graph.srcdata["h"] = feat_src
graph.update_all(fn.copy_u("h", "m"), fn.mean("m", "neigh"))
h_neigh = graph.dstdata["neigh"]
elif self._aggre_type == "gcn":
check_eq_shape(feat)
graph.srcdata["h"] = feat_src
graph.dstdata["h"] = feat_dst # same as above if homogeneous
graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "neigh"))
# divide in_degrees
degs = tf.cast(graph.in_degrees(), tf.float32)
h_neigh = (graph.dstdata["neigh"] + graph.dstdata["h"]) / (
tf.expand_dims(degs, -1) + 1
)
elif self._aggre_type == "pool":
graph.srcdata["h"] = tf.nn.relu(self.fc_pool(feat_src))
graph.update_all(fn.copy_u("h", "m"), fn.max("m", "neigh"))
h_neigh = graph.dstdata["neigh"]
elif self._aggre_type == "lstm":
graph.srcdata["h"] = feat_src
graph.update_all(fn.copy_u("h", "m"), self._lstm_reducer)
h_neigh = graph.dstdata["neigh"]
else:
raise KeyError(
"Aggregator type {} not recognized.".format(
self._aggre_type
)
)
# GraphSAGE GCN does not require fc_self.
if self._aggre_type == "gcn":
rst = self.fc_neigh(h_neigh)
else:
rst = self.fc_self(h_self) + self.fc_neigh(h_neigh)
# activation
if self.activation is not None:
rst = self.activation(rst)
# normalization
if self.norm is not None:
rst = self.norm(rst)
return rst
python/dgl/nn/tensorflow/conv/sgconv.py
+184
View File
@@ -0,0 +1,184 @@
"""tf Module for Simplifying Graph Convolution layer"""
# pylint: disable= no-member, arguments-differ, invalid-name, W0613
import numpy as np
import tensorflow as tf
from tensorflow.keras import layers
from .... import function as fn
from ....base import DGLError
class SGConv(layers.Layer):
r"""SGC layer from `Simplifying Graph
Convolutional Networks <https://arxiv.org/pdf/1902.07153.pdf>`__
.. math::
H^{K} = (\tilde{D}^{-1/2} \tilde{A} \tilde{D}^{-1/2})^K X \Theta
where :math:`\tilde{A}` is :math:`A` + :math:`I`.
Thus the graph input is expected to have self-loop edges added.
Parameters
----------
in_feats : int
Number of input features; i.e, the number of dimensions of :math:`X`.
out_feats : int
Number of output features; i.e, the number of dimensions of :math:`H^{K}`.
k : int
Number of hops :math:`K`. Defaults:``1``.
cached : bool
If True, the module would cache
.. math::
(\tilde{D}^{-\frac{1}{2}}\tilde{A}\tilde{D}^{-\frac{1}{2}})^K X\Theta
at the first forward call. This parameter should only be set to
``True`` in Transductive Learning setting.
bias : bool
If True, adds a learnable bias to the output. Default: ``True``.
norm : callable activation function/layer or None, optional
If not None, applies normalization to the updated node features. Default: ``False``.
allow_zero_in_degree : bool, optional
If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
since no message will be passed to those nodes. This is harmful for some applications
causing silent performance regression. This module will raise a DGLError if it detects
0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
and let the users handle it by themselves. Default: ``False``.
Note
----
Zero in-degree nodes will lead to invalid output value. This is because no message
will be passed to those nodes, the aggregation function will be appied on empty input.
A common practice to avoid this is to add a self-loop for each node in the graph if
it is homogeneous, which can be achieved by:
>>> g = ... # a DGLGraph
>>> g = dgl.add_self_loop(g)
Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
A common practise to handle this is to filter out the nodes with zero-in-degree when use
after conv.
Example
-------
>>> import dgl
>>> import numpy as np
>>> import tensorflow as tf
>>> from dgl.nn import SGConv
>>>
>>> with tf.device("CPU:0"):
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> g = dgl.add_self_loop(g)
>>> feat = tf.ones((6, 10))
>>> conv = SGConv(10, 2, k=2, cached=True)
>>> res = conv(g, feat)
>>> res
<tf.Tensor: shape=(6, 2), dtype=float32, numpy=
array([[0.61023676, 0.5246612 ],
[0.61023676, 0.5246612 ],
[0.61023676, 0.5246612 ],
[0.8697353 , 0.7477695 ],
[0.60570633, 0.520766 ],
[0.6102368 , 0.52466124]], dtype=float32)>
"""
def __init__(
self,
in_feats,
out_feats,
k=1,
cached=False,
bias=True,
norm=None,
allow_zero_in_degree=False,
):
super(SGConv, self).__init__()
self.fc = layers.Dense(out_feats, use_bias=bias)
self._cached = cached
self._cached_h = None
self._k = k
self.norm = norm
self._allow_zero_in_degree = allow_zero_in_degree
def set_allow_zero_in_degree(self, set_value):
r"""Set allow_zero_in_degree flag.
Parameters
----------
set_value : bool
The value to be set to the flag.
"""
self._allow_zero_in_degree = set_value
def call(self, graph, feat):
r"""Compute Simplifying Graph Convolution layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor
The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
is size of input feature, :math:`N` is the number of nodes.
Returns
-------
tf.Tensor
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
is size of output feature.
Raises
------
DGLError
If there are 0-in-degree nodes in the input graph, it will raise DGLError
since no message will be passed to those nodes. This will cause invalid output.
The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
Note
----
If ``cache`` is set to True, ``feat`` and ``graph`` should not change during
training, or you will get wrong results.
"""
with graph.local_scope():
if not self._allow_zero_in_degree:
if tf.math.count_nonzero(graph.in_degrees() == 0) > 0:
raise DGLError(
"There are 0-in-degree nodes in the graph, "
"output for those nodes will be invalid. "
"This is harmful for some applications, "
"causing silent performance regression. "
"Adding self-loop on the input graph by "
"calling `g = dgl.add_self_loop(g)` will resolve "
"the issue. Setting ``allow_zero_in_degree`` "
"to be `True` when constructing this module will "
"suppress the check and let the code run."
)
if self._cached_h is not None:
feat = self._cached_h
else:
# compute normalization
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)
norm = tf.expand_dims(norm, 1)
# compute (D^-1 A^k D)^k X
for _ in range(self._k):
feat = feat * norm
graph.ndata["h"] = feat
graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
feat = graph.ndata.pop("h")
feat = feat * norm
if self.norm is not None:
feat = self.norm(feat)
# cache feature
if self._cached:
self._cached_h = feat
return self.fc(feat)
python/dgl/nn/tensorflow/glob.py
+261
View File
@@ -0,0 +1,261 @@
"""Tensorflow modules for graph global pooling."""
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
import tensorflow as tf
from tensorflow.keras import layers
from ...readout import (
max_nodes,
mean_nodes,
softmax_nodes,
sum_nodes,
topk_nodes,
)
__all__ = [
"SumPooling",
"AvgPooling",
"MaxPooling",
"SortPooling",
"WeightAndSum",
"GlobalAttentionPooling",
]
class SumPooling(layers.Layer):
r"""Apply sum pooling over the nodes in the graph.
.. math::
r^{(i)} = \sum_{k=1}^{N_i} x^{(i)}_k
"""
def __init__(self):
super(SumPooling, self).__init__()
def call(self, graph, feat):
r"""Compute sum pooling.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor
The input feature with shape :math:`(N, *)` where
:math:`N` is the number of nodes in the graph.
Returns
-------
tf.Tensor
The output feature with shape :math:`(B, *)`, where
:math:`B` refers to the batch size.
"""
with graph.local_scope():
graph.ndata["h"] = feat
readout = sum_nodes(graph, "h")
return readout
class AvgPooling(layers.Layer):
r"""Apply average pooling over the nodes in the graph.
.. math::
r^{(i)} = \frac{1}{N_i}\sum_{k=1}^{N_i} x^{(i)}_k
"""
def __init__(self):
super(AvgPooling, self).__init__()
def call(self, graph, feat):
r"""Compute average pooling.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor
The input feature with shape :math:`(N, *)` where
:math:`N` is the number of nodes in the graph.
Returns
-------
tf.Tensor
The output feature with shape :math:`(B, *)`, where
:math:`B` refers to the batch size.
"""
with graph.local_scope():
graph.ndata["h"] = feat
readout = mean_nodes(graph, "h")
return readout
class MaxPooling(layers.Layer):
r"""Apply max pooling over the nodes in the graph.
.. math::
r^{(i)} = \max_{k=1}^{N_i}\left( x^{(i)}_k \right)
"""
def __init__(self):
super(MaxPooling, self).__init__()
def call(self, graph, feat):
r"""Compute max pooling.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor
The input feature with shape :math:`(N, *)` where
:math:`N` is the number of nodes in the graph.
Returns
-------
tf.Tensor
The output feature with shape :math:`(B, *)`, where
:math:`B` refers to the batch size.
"""
with graph.local_scope():
graph.ndata["h"] = feat
readout = max_nodes(graph, "h")
return readout
class SortPooling(layers.Layer):
r"""Sort Pooling from `An End-to-End Deep Learning Architecture for Graph Classification
<https://www.cse.wustl.edu/~ychen/public/DGCNN.pdf>`__
Parameters
----------
k : int
The number of nodes to hold for each graph.
"""
def __init__(self, k):
super(SortPooling, self).__init__()
self.k = k
def call(self, graph, feat):
r"""Compute sort pooling.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor
The input node feature with shape :math:`(N, D)` where
:math:`N` is the number of nodes in the graph.
Returns
-------
tf.Tensor
The output feature with shape :math:`(B, k * D)`, where
:math:`B` refers to the batch size.
"""
with graph.local_scope():
# Sort the feature of each node in ascending order.
feat = tf.sort(feat, -1)
graph.ndata["h"] = feat
# Sort nodes according to their last features.
ret = tf.reshape(
topk_nodes(graph, "h", self.k, sortby=-1)[0],
(-1, self.k * feat.shape[-1]),
)
return ret
class GlobalAttentionPooling(layers.Layer):
r"""Global Attention Pooling from `Gated Graph Sequence Neural Networks
<https://arxiv.org/abs/1511.05493.pdf>`__
.. math::
r^{(i)} = \sum_{k=1}^{N_i}\mathrm{softmax}\left(f_{gate}
\left(x^{(i)}_k\right)\right) f_{feat}\left(x^{(i)}_k\right)
Parameters
----------
gate_nn : tf.layers.Layer
A neural network that computes attention scores for each feature.
feat_nn : tf.layers.Layer, optional
A neural network applied to each feature before combining them
with attention scores.
"""
def __init__(self, gate_nn, feat_nn=None):
super(GlobalAttentionPooling, self).__init__()
self.gate_nn = gate_nn
self.feat_nn = feat_nn
def call(self, graph, feat):
r"""Compute global attention pooling.
Parameters
----------
graph : DGLGraph
The graph.
feat : tf.Tensor
The input node feature with shape :math:`(N, D)` where
:math:`N` is the number of nodes in the graph.
Returns
-------
tf.Tensor
The output feature with shape :math:`(B, *)`, where
:math:`B` refers to the batch size.
"""
with graph.local_scope():
gate = self.gate_nn(feat)
assert (
gate.shape[-1] == 1
), "The output of gate_nn should have size 1 at the last axis."
feat = self.feat_nn(feat) if self.feat_nn else feat
graph.ndata["gate"] = gate
gate = softmax_nodes(graph, "gate")
graph.ndata.pop("gate")
graph.ndata["r"] = feat * gate
readout = sum_nodes(graph, "r")
graph.ndata.pop("r")
return readout
class WeightAndSum(layers.Layer):
"""Compute importance weights for atoms and perform a weighted sum.
Parameters
----------
in_feats : int
Input atom feature size
"""
def __init__(self, in_feats):
super(WeightAndSum, self).__init__()
self.in_feats = in_feats
self.atom_weighting = tf.keras.Sequential(
layers.Dense(1), layers.Activation(tf.nn.sigmoid)
)
def call(self, g, feats):
"""Compute molecule representations out of atom representations
Parameters
----------
g : DGLGraph
DGLGraph with batch size B for processing multiple molecules in parallel
feats : FloatTensor of shape (N, self.in_feats)
Representations for all atoms in the molecules
* N is the total number of atoms in all molecules
Returns
-------
FloatTensor of shape (B, self.in_feats)
Representations for B molecules
"""
with g.local_scope():
g.ndata["h"] = feats
g.ndata["w"] = self.atom_weighting(g.ndata["h"])
h_g_sum = sum_nodes(g, "h", "w")
return h_g_sum
python/dgl/nn/tensorflow/hetero.py
+251
View File
@@ -0,0 +1,251 @@
"""Heterograph NN modules"""
import tensorflow as tf
from tensorflow.keras import layers
__all__ = ["HeteroGraphConv"]
class HeteroGraphConv(layers.Layer):
r"""A generic module for computing convolution on heterogeneous graphs.
The heterograph convolution applies sub-modules on their associating
relation graphs, which reads the features from source nodes and writes the
updated ones to destination nodes. If multiple relations have the same
destination node types, their results are aggregated by the specified method.
If the relation graph has no edge, the corresponding module will not be called.
Pseudo-code:
.. code::
outputs = {nty : [] for nty in g.dsttypes}
# Apply sub-modules on their associating relation graphs in parallel
for relation in g.canonical_etypes:
stype, etype, dtype = relation
dstdata = relation_submodule(g[relation], ...)
outputs[dtype].append(dstdata)
# Aggregate the results for each destination node type
rsts = {}
for ntype, ntype_outputs in outputs.items():
if len(ntype_outputs) != 0:
rsts[ntype] = aggregate(ntype_outputs)
return rsts
Examples
--------
Create a heterograph with three types of relations and nodes.
>>> import dgl
>>> g = dgl.heterograph({
... ('user', 'follows', 'user') : edges1,
... ('user', 'plays', 'game') : edges2,
... ('store', 'sells', 'game') : edges3})
Create a ``HeteroGraphConv`` that applies different convolution modules to
different relations. Note that the modules for ``'follows'`` and ``'plays'``
do not share weights.
>>> import dgl.nn.pytorch as dglnn
>>> conv = dglnn.HeteroGraphConv({
... 'follows' : dglnn.GraphConv(...),
... 'plays' : dglnn.GraphConv(...),
... 'sells' : dglnn.SAGEConv(...)},
... aggregate='sum')
Call forward with some ``'user'`` features. This computes new features for both
``'user'`` and ``'game'`` nodes.
>>> import tensorflow as tf
>>> h1 = {'user' : tf.random.normal((g.num_nodes('user'), 5))}
>>> h2 = conv(g, h1)
>>> print(h2.keys())
dict_keys(['user', 'game'])
Call forward with both ``'user'`` and ``'store'`` features. Because both the
``'plays'`` and ``'sells'`` relations will update the ``'game'`` features,
their results are aggregated by the specified method (i.e., summation here).
>>> f1 = {'user' : ..., 'store' : ...}
>>> f2 = conv(g, f1)
>>> print(f2.keys())
dict_keys(['user', 'game'])
Call forward with some ``'store'`` features. This only computes new features
for ``'game'`` nodes.
>>> g1 = {'store' : ...}
>>> g2 = conv(g, g1)
>>> print(g2.keys())
dict_keys(['game'])
Call forward with a pair of inputs is allowed and each submodule will also
be invoked with a pair of inputs.
>>> x_src = {'user' : ..., 'store' : ...}
>>> x_dst = {'user' : ..., 'game' : ...}
>>> y_dst = conv(g, (x_src, x_dst))
>>> print(y_dst.keys())
dict_keys(['user', 'game'])
Notes
-----
HeteroGraphConv requires that there is a module for every ``'etype'`` in an input graph.
If you want to apply HeteroGraphConv to a subset of a graph's ``'etypes'``, you must
create a new graph using for example :func:`~dgl.edge_type_subgraph()`.
Parameters
----------
mods : dict[str, nn.Module]
Modules associated with every edge types. The forward function of each
module must have a `DGLGraph` object as the first argument, and
its second argument is either a tensor object representing the node
features or a pair of tensor object representing the source and destination
node features.
aggregate : str, callable, optional
Method for aggregating node features generated by different relations.
Allowed string values are 'sum', 'max', 'min', 'mean', 'stack'.
The 'stack' aggregation is performed along the second dimension, whose order
is deterministic.
User can also customize the aggregator by providing a callable instance.
For example, aggregation by summation is equivalent to the follows:
.. code::
def my_agg_func(tensors, dsttype):
# tensors: is a list of tensors to aggregate
# dsttype: string name of the destination node type for which the
# aggregation is performed
stacked = tf.stack(tensors, axis=0)
return tf.reduce_sum(stacked, axis=0)
Attributes
----------
mods : dict[str, nn.Module]
Modules associated with every edge types.
"""
def __init__(self, mods, aggregate="sum"):
super(HeteroGraphConv, self).__init__()
self.mods = mods
# Do not break if graph has 0-in-degree nodes.
# Because there is no general rule to add self-loop for heterograph.
for _, v in self.mods.items():
set_allow_zero_in_degree_fn = getattr(
v, "set_allow_zero_in_degree", None
)
if callable(set_allow_zero_in_degree_fn):
set_allow_zero_in_degree_fn(True)
if isinstance(aggregate, str):
self.agg_fn = get_aggregate_fn(aggregate)
else:
self.agg_fn = aggregate
def call(self, g, inputs, mod_args=None, mod_kwargs=None):
"""Forward computation
Invoke the forward function with each module and aggregate their results.
Parameters
----------
g : DGLGraph
Graph data.
inputs : dict[str, Tensor] or pair of dict[str, Tensor]
Input node features.
mod_args : dict[str, tuple[any]], optional
Extra positional arguments for the sub-modules.
mod_kwargs : dict[str, dict[str, any]], optional
Extra key-word arguments for the sub-modules.
Returns
-------
dict[str, Tensor]
Output representations for every types of nodes.
"""
if mod_args is None:
mod_args = {}
if mod_kwargs is None:
mod_kwargs = {}
outputs = {nty: [] for nty in g.dsttypes}
if isinstance(inputs, tuple):
src_inputs, dst_inputs = inputs
for stype, etype, dtype in g.canonical_etypes:
rel_graph = g[stype, etype, dtype]
if stype not in src_inputs or dtype not in dst_inputs:
continue
dstdata = self.mods[etype](
rel_graph,
(src_inputs[stype], dst_inputs[dtype]),
*mod_args.get(etype, ()),
**mod_kwargs.get(etype, {})
)
outputs[dtype].append(dstdata)
else:
for stype, etype, dtype in g.canonical_etypes:
rel_graph = g[stype, etype, dtype]
if stype not in inputs:
continue
dstdata = self.mods[etype](
rel_graph,
(inputs[stype], inputs[dtype]),
*mod_args.get(etype, ()),
**mod_kwargs.get(etype, {})
)
outputs[dtype].append(dstdata)
rsts = {}
for nty, alist in outputs.items():
if len(alist) != 0:
rsts[nty] = self.agg_fn(alist, nty)
return rsts
def get_aggregate_fn(agg):
"""Internal function to get the aggregation function for node data
generated from different relations.
Parameters
----------
agg : str
Method for aggregating node features generated by different relations.
Allowed values are 'sum', 'max', 'min', 'mean', 'stack'.
Returns
-------
callable
Aggregator function that takes a list of tensors to aggregate
and returns one aggregated tensor.
"""
if agg == "sum":
fn = tf.reduce_sum
elif agg == "max":
fn = tf.reduce_max
elif agg == "min":
fn = tf.reduce_min
elif agg == "mean":
fn = tf.reduce_mean
elif agg == "stack":
fn = None # will not be called
else:
raise DGLError(
"Invalid cross type aggregator. Must be one of "
'"sum", "max", "min", "mean" or "stack". But got "%s"' % agg
)
if agg == "stack":
def stack_agg(inputs, dsttype): # pylint: disable=unused-argument
if len(inputs) == 0:
return None
return tf.stack(inputs, axis=1)
return stack_agg
else:
def aggfn(inputs, dsttype): # pylint: disable=unused-argument
if len(inputs) == 0:
return None
stacked = tf.stack(inputs, axis=0)
return fn(stacked, axis=0)
return aggfn
python/dgl/nn/tensorflow/softmax.py
+3
View File
@@ -0,0 +1,3 @@
"""tf modules for graph related softmax."""
# pylint: disable= unused-import
from ..functional import edge_softmax
python/dgl/nn/tensorflow/utils.py
+98
View File
@@ -0,0 +1,98 @@
"""Utilities for tf NN package"""
# pylint: disable=no-member, invalid-name
import tensorflow as tf
from tensorflow.keras import layers # pylint: disable=W0235
def matmul_maybe_select(A, B):
"""Perform Matrix multiplication C = A * B but A could be an integer id vector.
If A is an integer vector, we treat it as multiplying a one-hot encoded tensor.
In this case, the expensive dense matrix multiply can be replaced by a much
cheaper index lookup.
For example,
::
A = [2, 0, 1],
B = [[0.1, 0.2],
[0.3, 0.4],
[0.5, 0.6]]
then matmul_maybe_select(A, B) is equivalent to
::
[[0, 0, 1], [[0.1, 0.2],
[1, 0, 0], * [0.3, 0.4],
[0, 1, 0]] [0.5, 0.6]]
In all other cases, perform a normal matmul.
Parameters
----------
A : tf.Tensor
lhs tensor
B : tf.Tensor
rhs tensor
Returns
-------
C : tf.Tensor
result tensor
"""
if A.dtype == tf.int64 and len(A.shape) == 1:
return tf.gather(B, A)
else:
return tf.matmul(A, B)
def bmm_maybe_select(A, B, index):
"""Slice submatrices of A by the given index and perform bmm.
B is a 3D tensor of shape (N, D1, D2), which can be viewed as a stack of
N matrices of shape (D1, D2). The input index is an integer vector of length M.
A could be either:
(1) a dense tensor of shape (M, D1),
(2) an integer vector of length M.
The result C is a 2D matrix of shape (M, D2)
For case (1), C is computed by bmm:
::
C[i, :] = matmul(A[i, :], B[index[i], :, :])
For case (2), C is computed by index select:
::
C[i, :] = B[index[i], A[i], :]
Parameters
----------
A : tf.Tensor
lhs tensor
B : tf.Tensor
rhs tensor
index : tf.Tensor
index tensor
Returns
-------
C : tf.Tensor
return tensor
"""
if A.dtype == tf.int64 and len(A.shape) == 1:
# following is a faster version of B[index, A, :]
B = tf.reshape(B, (-1, B.shape[2]))
flatidx = index * B.shape[1] + A
return tf.gather(B, flatidx)
else:
BB = tf.gather(B, index)
return tf.squeeze(tf.matmul(tf.expand_dims(A, 1), BB), 1)
class Identity(layers.Layer):
"""A placeholder identity operator that is argument-insensitive."""
def call(self, x):
"""Return input"""
return x