chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,6 @@
|
||||
"""Package for mxnet-specific NN modules."""
|
||||
from .conv import *
|
||||
from .glob import *
|
||||
from .hetero import *
|
||||
from .softmax import *
|
||||
from .utils import Sequential
|
||||
@@ -0,0 +1,40 @@
|
||||
"""MXNet modules for graph convolutions."""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
|
||||
from .agnnconv import AGNNConv
|
||||
from .appnpconv import APPNPConv
|
||||
from .chebconv import ChebConv
|
||||
from .densechebconv import DenseChebConv
|
||||
from .densegraphconv import DenseGraphConv
|
||||
from .densesageconv import DenseSAGEConv
|
||||
from .edgeconv import EdgeConv
|
||||
from .gatconv import GATConv
|
||||
from .gatedgraphconv import GatedGraphConv
|
||||
from .ginconv import GINConv
|
||||
from .gmmconv import GMMConv
|
||||
from .graphconv import GraphConv
|
||||
from .nnconv import NNConv
|
||||
from .relgraphconv import RelGraphConv
|
||||
from .sageconv import SAGEConv
|
||||
from .sgconv import SGConv
|
||||
from .tagconv import TAGConv
|
||||
|
||||
__all__ = [
|
||||
"GraphConv",
|
||||
"TAGConv",
|
||||
"RelGraphConv",
|
||||
"GATConv",
|
||||
"SAGEConv",
|
||||
"GatedGraphConv",
|
||||
"ChebConv",
|
||||
"AGNNConv",
|
||||
"APPNPConv",
|
||||
"DenseGraphConv",
|
||||
"DenseSAGEConv",
|
||||
"DenseChebConv",
|
||||
"EdgeConv",
|
||||
"GINConv",
|
||||
"GMMConv",
|
||||
"NNConv",
|
||||
"SGConv",
|
||||
]
|
||||
@@ -0,0 +1,163 @@
|
||||
"""MXNet Module for Attention-based Graph Neural Network layer"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import mxnet as mx
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from .... import function as fn
|
||||
from ....base import DGLError
|
||||
from ....utils import expand_as_pair
|
||||
from ...functional import edge_softmax
|
||||
from ..utils import normalize
|
||||
|
||||
|
||||
class AGNNConv(nn.Block):
|
||||
r"""Attention-based Graph Neural Network layer from `Attention-based Graph Neural Network for
|
||||
Semi-Supervised Learning <https://arxiv.org/abs/1803.03735>`__
|
||||
|
||||
.. math::
|
||||
H^{l+1} = P H^{l}
|
||||
|
||||
where :math:`P` is computed as:
|
||||
|
||||
.. math::
|
||||
P_{ij} = \mathrm{softmax}_i ( \beta \cdot \cos(h_i^l, h_j^l))
|
||||
|
||||
where :math:`\beta` is a single scalar parameter.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
init_beta : float, optional
|
||||
The :math:`\beta` in the formula, a single scalar parameter.
|
||||
learn_beta : bool, optional
|
||||
If True, :math:`\beta` will be learnable parameter.
|
||||
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 mxnet as mx
|
||||
>>> from dgl.nn import AGNNConv
|
||||
>>>
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> g = dgl.add_self_loop(g)
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = AGNNConv()
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat)
|
||||
>>> res
|
||||
[[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
|
||||
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
|
||||
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
|
||||
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
|
||||
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
|
||||
[1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]]
|
||||
<NDArray 6x10 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, init_beta=1.0, learn_beta=True, allow_zero_in_degree=False
|
||||
):
|
||||
super(AGNNConv, self).__init__()
|
||||
self._allow_zero_in_degree = allow_zero_in_degree
|
||||
with self.name_scope():
|
||||
self.beta = self.params.get(
|
||||
"beta",
|
||||
shape=(1,),
|
||||
grad_req="write" if learn_beta else "null",
|
||||
init=mx.init.Constant(init_beta),
|
||||
)
|
||||
|
||||
def set_allow_zero_in_degree(self, set_value):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
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 forward(self, graph, feat):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute AGNN layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
The input feature of shape :math:`(N, *)` :math:`N` is the
|
||||
number of nodes, and :math:`*` could be of any shape.
|
||||
If a pair of mxnet.NDArray is given, the pair must contain two tensors of shape
|
||||
:math:`(N_{in}, *)` and :math:`(N_{out}, *)`, the :math:`*` in the later
|
||||
tensor must equal the previous one.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature of shape :math:`(N, *)` where :math:`*`
|
||||
should be the same as input shape.
|
||||
|
||||
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 graph.in_degrees().min() == 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)
|
||||
graph.srcdata["h"] = feat_src
|
||||
graph.srcdata["norm_h"] = normalize(feat_src, p=2, axis=-1)
|
||||
if isinstance(feat, tuple) or graph.is_block:
|
||||
graph.dstdata["norm_h"] = normalize(feat_dst, p=2, axis=-1)
|
||||
# compute cosine distance
|
||||
graph.apply_edges(fn.u_dot_v("norm_h", "norm_h", "cos"))
|
||||
cos = graph.edata.pop("cos")
|
||||
e = self.beta.data(feat_src.context) * cos
|
||||
graph.edata["p"] = edge_softmax(graph, e)
|
||||
graph.update_all(fn.u_mul_e("h", "p", "m"), fn.sum("m", "h"))
|
||||
return graph.dstdata.pop("h")
|
||||
@@ -0,0 +1,113 @@
|
||||
"""MXNet Module for APPNPConv"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import mxnet as mx
|
||||
from mxnet import nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from .... import function as fn
|
||||
|
||||
|
||||
class APPNPConv(nn.Block):
|
||||
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^{l+1} &= (1-\alpha)\left(\tilde{D}^{-1/2}
|
||||
\tilde{A} \tilde{D}^{-1/2} H^{l}\right) + \alpha H^{0}
|
||||
|
||||
where :math:`\tilde{A}` is :math:`A` + :math:`I`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
k : int
|
||||
The number of iterations :math:`K`.
|
||||
alpha : float
|
||||
The teleport probability :math:`\alpha`.
|
||||
edge_drop : float, optional
|
||||
The dropout rate on edges that controls the
|
||||
messages received by each node. Default: ``0``.
|
||||
|
||||
Example
|
||||
-------
|
||||
>>> import dgl
|
||||
>>> import numpy as np
|
||||
>>> import mxnet as mx
|
||||
>>> from dgl.nn import APPNPConv
|
||||
>>>
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = APPNPConv(k=3, alpha=0.5)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat)
|
||||
>>> res
|
||||
[[1. 1. 1. 1. 1. 1.
|
||||
1. 1. 1. 1. ]
|
||||
[1. 1. 1. 1. 1. 1.
|
||||
1. 1. 1. 1. ]
|
||||
[1. 1. 1. 1. 1. 1.
|
||||
1. 1. 1. 1. ]
|
||||
[1.0303301 1.0303301 1.0303301 1.0303301 1.0303301 1.0303301
|
||||
1.0303301 1.0303301 1.0303301 1.0303301 ]
|
||||
[0.86427665 0.86427665 0.86427665 0.86427665 0.86427665 0.86427665
|
||||
0.86427665 0.86427665 0.86427665 0.86427665]
|
||||
[0.5 0.5 0.5 0.5 0.5 0.5
|
||||
0.5 0.5 0.5 0.5 ]]
|
||||
<NDArray 6x10 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(self, k, alpha, edge_drop=0.0):
|
||||
super(APPNPConv, self).__init__()
|
||||
self._k = k
|
||||
self._alpha = alpha
|
||||
with self.name_scope():
|
||||
self.edge_drop = nn.Dropout(edge_drop)
|
||||
|
||||
def forward(self, graph, feat):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute APPNP layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mx.NDArray
|
||||
The input feature of shape :math:`(N, *)`. :math:`N` is the
|
||||
number of nodes, and :math:`*` could be of any shape.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mx.NDArray
|
||||
The output feature of shape :math:`(N, *)` where :math:`*`
|
||||
should be the same as input shape.
|
||||
"""
|
||||
with graph.local_scope():
|
||||
norm = mx.nd.power(
|
||||
mx.nd.clip(
|
||||
graph.in_degrees().astype(feat.dtype),
|
||||
a_min=1,
|
||||
a_max=float("inf"),
|
||||
),
|
||||
-0.5,
|
||||
)
|
||||
shp = norm.shape + (1,) * (feat.ndim - 1)
|
||||
norm = norm.reshape(shp).as_in_context(feat.context)
|
||||
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(
|
||||
nd.ones((graph.num_edges(), 1), ctx=feat.context)
|
||||
)
|
||||
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
|
||||
@@ -0,0 +1,167 @@
|
||||
"""MXNet Module for Chebyshev Spectral Graph Convolution layer"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet import nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from .... import broadcast_nodes, function as fn
|
||||
from ....base import dgl_warning
|
||||
|
||||
|
||||
class ChebConv(nn.Block):
|
||||
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 mxnet as mx
|
||||
>>> from dgl.nn import ChebConv
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = ChebConv(10, 2, 2)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat)
|
||||
>>> res
|
||||
[[ 0.832592 -0.738757 ]
|
||||
[ 0.832592 -0.738757 ]
|
||||
[ 0.832592 -0.738757 ]
|
||||
[ 0.43377423 -1.0455742 ]
|
||||
[ 1.1145986 -0.5218046 ]
|
||||
[ 1.7954229 0.00196505]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(self, in_feats, out_feats, k, bias=True):
|
||||
super(ChebConv, self).__init__()
|
||||
self._in_feats = in_feats
|
||||
self._out_feats = out_feats
|
||||
self._k = k
|
||||
with self.name_scope():
|
||||
self.fc = nn.Sequential()
|
||||
for _ in range(k):
|
||||
self.fc.add(
|
||||
nn.Dense(
|
||||
out_feats,
|
||||
use_bias=False,
|
||||
weight_initializer=mx.init.Xavier(
|
||||
magnitude=math.sqrt(2.0)
|
||||
),
|
||||
in_units=in_feats,
|
||||
)
|
||||
)
|
||||
if bias:
|
||||
self.bias = self.params.get(
|
||||
"bias", shape=(out_feats,), init=mx.init.Zero()
|
||||
)
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
def forward(self, graph, feat, lambda_max=None):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute ChebNet layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
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
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
|
||||
is size of output feature.
|
||||
"""
|
||||
with graph.local_scope():
|
||||
degs = graph.in_degrees().astype("float32")
|
||||
norm = mx.nd.power(
|
||||
mx.nd.clip(degs, a_min=1, a_max=float("inf")), -0.5
|
||||
)
|
||||
norm = norm.expand_dims(-1).as_in_context(feat.context)
|
||||
|
||||
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 = nd.array(lambda_max).as_in_context(feat.context)
|
||||
if lambda_max.ndim == 1:
|
||||
lambda_max = lambda_max.expand_dims(-1)
|
||||
# broadcast from (B, 1) to (N, 1)
|
||||
lambda_max = broadcast_nodes(graph, lambda_max)
|
||||
# T0(X)
|
||||
Tx_0 = feat
|
||||
rst = self.fc[0](Tx_0)
|
||||
# T1(X)
|
||||
if self._k > 1:
|
||||
graph.ndata["h"] = Tx_0 * norm
|
||||
graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
|
||||
h = graph.ndata.pop("h") * norm
|
||||
# Λ = 2 * (I - D ^ -1/2 A D ^ -1/2) / lambda_max - I
|
||||
# = - 2(D ^ -1/2 A D ^ -1/2) / lambda_max + (2 / lambda_max - 1) I
|
||||
Tx_1 = -2.0 * h / lambda_max + Tx_0 * (2.0 / lambda_max - 1)
|
||||
rst = rst + self.fc[1](Tx_1)
|
||||
# Ti(x), i = 2...k
|
||||
for i in range(2, self._k):
|
||||
graph.ndata["h"] = Tx_1 * norm
|
||||
graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
|
||||
h = graph.ndata.pop("h") * norm
|
||||
# Tx_k = 2 * Λ * Tx_(k-1) - Tx_(k-2)
|
||||
# = - 4(D ^ -1/2 A D ^ -1/2) / lambda_max Tx_(k-1) +
|
||||
# (4 / lambda_max - 2) Tx_(k-1) -
|
||||
# Tx_(k-2)
|
||||
Tx_2 = (
|
||||
-4.0 * h / lambda_max + Tx_1 * (4.0 / lambda_max - 2) - Tx_0
|
||||
)
|
||||
rst = rst + self.fc[i](Tx_2)
|
||||
Tx_1, Tx_0 = Tx_2, Tx_1
|
||||
# add bias
|
||||
if self.bias is not None:
|
||||
rst = rst + self.bias.data(feat.context)
|
||||
return rst
|
||||
@@ -0,0 +1,109 @@
|
||||
"""MXNet Module for DenseChebConv"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet import nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
|
||||
class DenseChebConv(nn.Block):
|
||||
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.pytorch.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
|
||||
with self.name_scope():
|
||||
self.fc = nn.Sequential()
|
||||
for _ in range(k):
|
||||
self.fc.add(
|
||||
nn.Dense(
|
||||
out_feats,
|
||||
in_units=in_feats,
|
||||
use_bias=False,
|
||||
weight_initializer=mx.init.Xavier(
|
||||
magnitude=math.sqrt(2.0)
|
||||
),
|
||||
)
|
||||
)
|
||||
if bias:
|
||||
self.bias = self.params.get(
|
||||
"bias", shape=(out_feats,), init=mx.init.Zero()
|
||||
)
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
def forward(self, adj, feat, lambda_max=None):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute (Dense) Chebyshev Spectral Graph Convolution layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
adj : mxnet.NDArray
|
||||
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 : mxnet.NDArray
|
||||
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
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
|
||||
is size of output feature.
|
||||
"""
|
||||
A = adj.astype(feat.dtype).as_in_context(feat.context)
|
||||
num_nodes = A.shape[0]
|
||||
|
||||
in_degree = 1.0 / nd.clip(A.sum(axis=1), 1, float("inf")).sqrt()
|
||||
D_invsqrt = nd.diag(in_degree)
|
||||
I = nd.eye(num_nodes, ctx=A.context)
|
||||
L = I - nd.dot(D_invsqrt, nd.dot(A, D_invsqrt))
|
||||
|
||||
if lambda_max is None:
|
||||
# NOTE(zihao): this only works for directed graph.
|
||||
lambda_max = (nd.linalg.syevd(L)[1]).max()
|
||||
|
||||
L_hat = 2 * L / lambda_max - I
|
||||
Z = [nd.eye(num_nodes, ctx=A.context)]
|
||||
Zh = self.fc[0](feat)
|
||||
for i in range(1, self._k):
|
||||
if i == 1:
|
||||
Z.append(L_hat)
|
||||
else:
|
||||
Z.append(2 * nd.dot(L_hat, Z[-1]) - Z[-2])
|
||||
Zh = Zh + nd.dot(Z[i], self.fc[i](feat))
|
||||
|
||||
if self.bias is not None:
|
||||
Zh = Zh + self.bias.data(feat.context)
|
||||
return Zh
|
||||
@@ -0,0 +1,126 @@
|
||||
"""MXNet Module for DenseGraphConv"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet import nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
|
||||
class DenseGraphConv(nn.Block):
|
||||
"""Graph Convolutional layer from `Semi-Supervised Classification with Graph
|
||||
Convolutional Networks <https://arxiv.org/abs/1609.02907>`__
|
||||
|
||||
We recommend user to use this module when applying graph convolution on
|
||||
dense graphs.
|
||||
|
||||
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. If is `'right'`, divide the aggregated messages
|
||||
by each node's in-degrees, which is equivalent to averaging the received messages.
|
||||
If is `'none'`, no normalization is applied. Default is `'both'`,
|
||||
where the :math:`c_{ij}` in the paper is applied.
|
||||
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``.
|
||||
|
||||
Notes
|
||||
-----
|
||||
Zero in-degree nodes will lead to all-zero output. A common practice
|
||||
to avoid this is to add a self-loop for each node in the graph,
|
||||
which can be achieved by setting the diagonal of the adjacency matrix to be 1.
|
||||
|
||||
See also
|
||||
--------
|
||||
`GraphConv <https://docs.dgl.ai/api/python/nn.pytorch.html#graphconv>`__
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, in_feats, out_feats, norm="both", bias=True, activation=None
|
||||
):
|
||||
super(DenseGraphConv, self).__init__()
|
||||
self._in_feats = in_feats
|
||||
self._out_feats = out_feats
|
||||
self._norm = norm
|
||||
with self.name_scope():
|
||||
self.weight = self.params.get(
|
||||
"weight",
|
||||
shape=(in_feats, out_feats),
|
||||
init=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
if bias:
|
||||
self.bias = self.params.get(
|
||||
"bias", shape=(out_feats,), init=mx.init.Zero()
|
||||
)
|
||||
else:
|
||||
self.bias = None
|
||||
self._activation = activation
|
||||
|
||||
def forward(self, adj, feat):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute (Dense) Graph Convolution layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
adj : mxnet.NDArray
|
||||
The adjacency matrix of the graph to apply Graph Convolution on, when
|
||||
applied to a unidirectional bipartite graph, ``adj`` should be of shape
|
||||
should be of shape :math:`(N_{out}, N_{in})`; when applied to a homo
|
||||
graph, ``adj`` should be of shape :math:`(N, N)`. In both cases,
|
||||
a row represents a destination node while a column represents a source
|
||||
node.
|
||||
feat : mxnet.NDArray
|
||||
The input feature.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
|
||||
is size of output feature.
|
||||
"""
|
||||
adj = adj.astype(feat.dtype).as_in_context(feat.context)
|
||||
src_degrees = nd.clip(adj.sum(axis=0), a_min=1, a_max=float("inf"))
|
||||
dst_degrees = nd.clip(adj.sum(axis=1), a_min=1, a_max=float("inf"))
|
||||
feat_src = feat
|
||||
|
||||
if self._norm == "both":
|
||||
norm_src = nd.power(src_degrees, -0.5)
|
||||
shp_src = norm_src.shape + (1,) * (feat.ndim - 1)
|
||||
norm_src = norm_src.reshape(shp_src).as_in_context(feat.context)
|
||||
feat_src = feat_src * norm_src
|
||||
|
||||
if self._in_feats > self._out_feats:
|
||||
# mult W first to reduce the feature size for aggregation.
|
||||
feat_src = nd.dot(feat_src, self.weight.data(feat_src.context))
|
||||
rst = nd.dot(adj, feat_src)
|
||||
else:
|
||||
# aggregate first then mult W
|
||||
rst = nd.dot(adj, feat_src)
|
||||
rst = nd.dot(rst, self.weight.data(feat_src.context))
|
||||
|
||||
if self._norm != "none":
|
||||
if self._norm == "both":
|
||||
norm_dst = nd.power(dst_degrees, -0.5)
|
||||
else: # right
|
||||
norm_dst = 1.0 / dst_degrees
|
||||
shp_dst = norm_dst.shape + (1,) * (feat.ndim - 1)
|
||||
norm_dst = norm_dst.reshape(shp_dst).as_in_context(feat.context)
|
||||
rst = rst * norm_dst
|
||||
|
||||
if self.bias is not None:
|
||||
rst = rst + self.bias.data(feat.context)
|
||||
|
||||
if self._activation is not None:
|
||||
rst = self._activation(rst)
|
||||
|
||||
return rst
|
||||
@@ -0,0 +1,109 @@
|
||||
"""MXNet Module for DenseGraphSAGE"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet import nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from ....utils import check_eq_shape
|
||||
|
||||
|
||||
class DenseSAGEConv(nn.Block):
|
||||
"""GraphSAGE layer from `Inductive Representation Learning on Large Graphs
|
||||
<https://arxiv.org/abs/1706.02216>`__
|
||||
|
||||
We recommend to use this module when appying GraphSAGE on dense graphs.
|
||||
|
||||
Note that we only support gcn aggregator in DenseSAGEConv.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
in_feats : int
|
||||
Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
|
||||
out_feats : int
|
||||
Output feature size; i.e, the number of dimensions of :math:`h_i^{(l+1)}`.
|
||||
feat_drop : float, optional
|
||||
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``.
|
||||
|
||||
See also
|
||||
--------
|
||||
`SAGEConv <https://docs.dgl.ai/api/python/nn.pytorch.html#sageconv>`__
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_feats,
|
||||
out_feats,
|
||||
feat_drop=0.0,
|
||||
bias=True,
|
||||
norm=None,
|
||||
activation=None,
|
||||
):
|
||||
super(DenseSAGEConv, self).__init__()
|
||||
self._in_feats = in_feats
|
||||
self._out_feats = out_feats
|
||||
self._norm = norm
|
||||
with self.name_scope():
|
||||
self.feat_drop = nn.Dropout(feat_drop)
|
||||
self.activation = activation
|
||||
self.fc = nn.Dense(
|
||||
out_feats,
|
||||
in_units=in_feats,
|
||||
use_bias=bias,
|
||||
weight_initializer=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
|
||||
def forward(self, adj, feat):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute (Dense) Graph SAGE layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
adj : mxnet.NDArray
|
||||
The adjacency matrix of the graph to apply SAGE Convolution on, when
|
||||
applied to a unidirectional bipartite graph, ``adj`` should be of shape
|
||||
should be of shape :math:`(N_{out}, N_{in})`; when applied to a homo
|
||||
graph, ``adj`` should be of shape :math:`(N, N)`. In both cases,
|
||||
a row represents a destination node while a column represents a source
|
||||
node.
|
||||
feat : mxnet.NDArray or a pair of mxnet.NDArray
|
||||
If a mxnet.NDArray 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 mxnet.NDArray is given, the pair must contain two tensors of shape
|
||||
:math:`(N_{in}, D_{in})` and :math:`(N_{out}, D_{in})`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
|
||||
is size of output feature.
|
||||
"""
|
||||
check_eq_shape(feat)
|
||||
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)
|
||||
adj = adj.astype(feat_src.dtype).as_in_context(feat_src.context)
|
||||
in_degrees = adj.sum(axis=1, keepdims=True)
|
||||
h_neigh = (nd.dot(adj, feat_src) + feat_dst) / (in_degrees + 1)
|
||||
rst = self.fc(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
|
||||
@@ -0,0 +1,192 @@
|
||||
"""MXNet Module for EdgeConv Layer"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import mxnet as mx
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from .... import function as fn
|
||||
from ....base import DGLError
|
||||
from ....utils import expand_as_pair
|
||||
|
||||
|
||||
class EdgeConv(nn.Block):
|
||||
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.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import dgl
|
||||
>>> import numpy as np
|
||||
>>> import mxnet as mx
|
||||
>>> from mxnet import gluon
|
||||
>>> from dgl.nn import EdgeConv
|
||||
>>>
|
||||
>>> # Case 1: Homogeneous graph
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> g = dgl.add_self_loop(g)
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = EdgeConv(10, 2)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat)
|
||||
>>> res
|
||||
[[1.0517545 0.8091326]
|
||||
[1.0517545 0.8091326]
|
||||
[1.0517545 0.8091326]
|
||||
[1.0517545 0.8091326]
|
||||
[1.0517545 0.8091326]
|
||||
[1.0517545 0.8091326]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
|
||||
>>> # Case 2: Unidirectional bipartite graph
|
||||
>>> u = [0, 1, 0, 0, 1]
|
||||
>>> v = [0, 1, 2, 3, 2]
|
||||
>>> g = dgl.bipartite((u, v))
|
||||
>>> u_fea = mx.nd.random.randn(2, 5)
|
||||
>>> v_fea = mx.nd.random.randn(4, 5)
|
||||
>>> conv = EdgeConv(5, 2, 3)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, (u_fea, v_fea))
|
||||
>>> res
|
||||
[[-3.4617817 0.84700686]
|
||||
[ 1.3170856 -1.5731761 ]
|
||||
[-2.0761423 0.56653017]
|
||||
[-1.015364 0.78919804]]
|
||||
<NDArray 4x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, in_feat, out_feat, 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
|
||||
|
||||
with self.name_scope():
|
||||
self.theta = nn.Dense(
|
||||
out_feat, in_units=in_feat, weight_initializer=mx.init.Xavier()
|
||||
)
|
||||
self.phi = nn.Dense(
|
||||
out_feat, in_units=in_feat, weight_initializer=mx.init.Xavier()
|
||||
)
|
||||
|
||||
if batch_norm:
|
||||
self.bn = nn.BatchNorm(in_channels=out_feat)
|
||||
|
||||
def set_allow_zero_in_degree(self, set_value):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
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 forward(self, g, h):
|
||||
"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Forward computation
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray or pair of mxnet.NDArray
|
||||
:math:`(N, D)` where :math:`N` is the number of nodes and
|
||||
:math:`D` is the number of feature dimensions.
|
||||
|
||||
If a pair of mxnet.NDArray 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
|
||||
-------
|
||||
mxnet.NDArray
|
||||
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 g.in_degrees().min() == 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(h, 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"))
|
||||
g.edata["e"] = self.bn(g.edata["e"])
|
||||
g.update_all(fn.copy_e("e", "m"), fn.max("m", "x"))
|
||||
return g.dstdata["x"]
|
||||
@@ -0,0 +1,341 @@
|
||||
"""MXNet modules for graph attention networks(GAT)."""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet.gluon import nn
|
||||
from mxnet.gluon.contrib.nn import Identity
|
||||
|
||||
from .... import function as fn
|
||||
from ....base import DGLError
|
||||
from ....utils import expand_as_pair
|
||||
from ...functional import edge_softmax
|
||||
|
||||
|
||||
# pylint: enable=W0235
|
||||
class GATConv(nn.Block):
|
||||
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)}`.
|
||||
GATConv 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 mxnet as mx
|
||||
>>> from mxnet import gluon
|
||||
>>> from dgl.nn import GATConv
|
||||
>>>
|
||||
>>> # Case 1: Homogeneous graph
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> g = dgl.add_self_loop(g)
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> gatconv = GATConv(10, 2, num_heads=3)
|
||||
>>> gatconv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = gatconv(g, feat)
|
||||
>>> res
|
||||
[[[ 0.32368395 -0.10501936]
|
||||
[ 1.0839728 0.92690575]
|
||||
[-0.54581136 -0.84279203]]
|
||||
[[ 0.32368395 -0.10501936]
|
||||
[ 1.0839728 0.92690575]
|
||||
[-0.54581136 -0.84279203]]
|
||||
[[ 0.32368395 -0.10501936]
|
||||
[ 1.0839728 0.92690575]
|
||||
[-0.54581136 -0.84279203]]
|
||||
[[ 0.32368395 -0.10501937]
|
||||
[ 1.0839728 0.9269058 ]
|
||||
[-0.5458114 -0.8427921 ]]
|
||||
[[ 0.32368395 -0.10501936]
|
||||
[ 1.0839728 0.92690575]
|
||||
[-0.54581136 -0.84279203]]
|
||||
[[ 0.32368395 -0.10501936]
|
||||
[ 1.0839728 0.92690575]
|
||||
[-0.54581136 -0.84279203]]]
|
||||
<NDArray 6x3x2 @cpu(0)>
|
||||
|
||||
>>> # 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)})
|
||||
>>> u_feat = mx.nd.random.randn(2, 5)
|
||||
>>> v_feat = mx.nd.random.randn(4, 10)
|
||||
>>> gatconv = GATConv((5,10), 2, 3)
|
||||
>>> gatconv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = gatconv(g, (u_feat, v_feat))
|
||||
>>> res
|
||||
[[[-1.01624 1.8138596 ]
|
||||
[ 1.2322129 -0.8410206 ]
|
||||
[-1.9325689 1.3824553 ]]
|
||||
[[ 0.9915016 -1.6564168 ]
|
||||
[-0.32610354 0.42505783]
|
||||
[ 1.5278397 -0.92114615]]
|
||||
[[-0.32592064 0.62067866]
|
||||
[ 0.6162219 -0.3405491 ]
|
||||
[-1.356375 0.9988818 ]]
|
||||
[[-1.01624 1.8138596 ]
|
||||
[ 1.2322129 -0.8410206 ]
|
||||
[-1.9325689 1.3824553 ]]]
|
||||
<NDArray 4x3x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
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_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
|
||||
self._in_feats = in_feats
|
||||
self._out_feats = out_feats
|
||||
self._allow_zero_in_degree = allow_zero_in_degree
|
||||
with self.name_scope():
|
||||
if isinstance(in_feats, tuple):
|
||||
self.fc_src = nn.Dense(
|
||||
out_feats * num_heads,
|
||||
use_bias=False,
|
||||
weight_initializer=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
in_units=self._in_src_feats,
|
||||
)
|
||||
self.fc_dst = nn.Dense(
|
||||
out_feats * num_heads,
|
||||
use_bias=False,
|
||||
weight_initializer=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
in_units=self._in_dst_feats,
|
||||
)
|
||||
else:
|
||||
self.fc = nn.Dense(
|
||||
out_feats * num_heads,
|
||||
use_bias=False,
|
||||
weight_initializer=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
in_units=in_feats,
|
||||
)
|
||||
self.attn_l = self.params.get(
|
||||
"attn_l",
|
||||
shape=(1, num_heads, out_feats),
|
||||
init=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
self.attn_r = self.params.get(
|
||||
"attn_r",
|
||||
shape=(1, num_heads, out_feats),
|
||||
init=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
self.feat_drop = nn.Dropout(feat_drop)
|
||||
self.attn_drop = nn.Dropout(attn_drop)
|
||||
self.leaky_relu = nn.LeakyReLU(negative_slope)
|
||||
if residual:
|
||||
if in_feats != out_feats:
|
||||
self.res_fc = nn.Dense(
|
||||
out_feats * num_heads,
|
||||
use_bias=False,
|
||||
weight_initializer=mx.init.Xavier(
|
||||
magnitude=math.sqrt(2.0)
|
||||
),
|
||||
in_units=in_feats,
|
||||
)
|
||||
else:
|
||||
self.res_fc = Identity()
|
||||
else:
|
||||
self.res_fc = None
|
||||
self.activation = activation
|
||||
|
||||
def set_allow_zero_in_degree(self, set_value):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
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 forward(self, graph, feat, get_attention=False):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute graph attention network layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray or pair of mxnet.NDArray
|
||||
If a mxnet.NDArray 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 mxnet.NDArray 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
|
||||
-------
|
||||
mxnet.NDArray
|
||||
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.
|
||||
mxnet.NDArray, 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 graph.in_degrees().min() == 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 = feat[0].shape[:-1]
|
||||
dst_prefix_shape = feat[1].shape[:-1]
|
||||
feat_dim = feat[0].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 = self.fc_src(h_src.reshape(-1, feat_dim)).reshape(
|
||||
*src_prefix_shape, self._num_heads, self._out_feats
|
||||
)
|
||||
feat_dst = self.fc_dst(h_dst.reshape(-1, feat_dim)).reshape(
|
||||
*dst_prefix_shape, self._num_heads, self._out_feats
|
||||
)
|
||||
else:
|
||||
src_prefix_shape = dst_prefix_shape = feat.shape[:-1]
|
||||
feat_dim = feat[0].shape[-1]
|
||||
h_src = h_dst = self.feat_drop(feat)
|
||||
feat_src = feat_dst = self.fc(
|
||||
h_src.reshape(-1, feat_dim)
|
||||
).reshape(*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 = (
|
||||
(feat_src * self.attn_l.data(feat_src.context))
|
||||
.sum(axis=-1)
|
||||
.expand_dims(-1)
|
||||
)
|
||||
er = (
|
||||
(feat_dst * self.attn_r.data(feat_src.context))
|
||||
.sum(axis=-1)
|
||||
.expand_dims(-1)
|
||||
)
|
||||
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))
|
||||
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 = self.res_fc(h_dst.reshape(-1, feat_dim)).reshape(
|
||||
*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
|
||||
@@ -0,0 +1,135 @@
|
||||
"""MXNet Module for Gated Graph Convolution layer"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name, cell-var-from-loop
|
||||
import mxnet as mx
|
||||
from mxnet import gluon, nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from .... import function as fn
|
||||
|
||||
|
||||
class GatedGraphConv(nn.Block):
|
||||
r"""Gated Graph Convolution layer from `Gated Graph Sequence
|
||||
Neural Networks <https://arxiv.org/pdf/1511.05493.pdf>`__
|
||||
|
||||
.. math::
|
||||
h_{i}^{0} &= [ x_i \| \mathbf{0} ]
|
||||
|
||||
a_{i}^{t} &= \sum_{j\in\mathcal{N}(i)} W_{e_{ij}} h_{j}^{t}
|
||||
|
||||
h_{i}^{t+1} &= \mathrm{GRU}(a_{i}^{t}, h_{i}^{t})
|
||||
|
||||
Parameters
|
||||
----------
|
||||
in_feats : int
|
||||
Input feature size; i.e, the number of dimensions of :math:`x_i`.
|
||||
out_feats : int
|
||||
Output feature size; i.e., the number of dimensions of :math:`h_i^{(t+1)}`.
|
||||
n_steps : int
|
||||
Number of recurrent steps; i.e, the :math:`t` in the above formula.
|
||||
n_etypes : int
|
||||
Number of edge types.
|
||||
bias : bool
|
||||
If True, adds a learnable bias to the output. Default: ``True``.
|
||||
Can only be set to True in MXNet.
|
||||
|
||||
Example
|
||||
-------
|
||||
>>> import dgl
|
||||
>>> import numpy as np
|
||||
>>> import mxnet as mx
|
||||
>>> from dgl.nn import GatedGraphConv
|
||||
>>>
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = GatedGraphConv(10, 10, 2, 3)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> etype = mx.nd.array([0,1,2,0,1,2])
|
||||
>>> res = conv(g, feat, etype)
|
||||
>>> res
|
||||
[[0.24378185 0.17402579 0.2644723 0.2740628 0.14041871 0.32523093
|
||||
0.2703067 0.18234392 0.32777587 0.30957845]
|
||||
[0.17872348 0.28878236 0.2509409 0.20139427 0.3355541 0.22643831
|
||||
0.2690711 0.22341749 0.27995753 0.21575949]
|
||||
[0.23911178 0.16696918 0.26120248 0.27397877 0.13745922 0.3223175
|
||||
0.27561218 0.18071817 0.3251124 0.30608907]
|
||||
[0.25242943 0.3098581 0.25249368 0.27968448 0.24624602 0.12270881
|
||||
0.335147 0.31550157 0.19065917 0.21087633]
|
||||
[0.17503153 0.29523152 0.2474858 0.20848347 0.3526433 0.23443702
|
||||
0.24741334 0.21986549 0.28935105 0.21859099]
|
||||
[0.2159364 0.26942077 0.23083271 0.28329757 0.24758333 0.24230732
|
||||
0.23958017 0.23430146 0.26431587 0.27001363]]
|
||||
<NDArray 6x10 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(self, in_feats, out_feats, n_steps, n_etypes, bias=True):
|
||||
super(GatedGraphConv, self).__init__()
|
||||
self._in_feats = in_feats
|
||||
self._out_feats = out_feats
|
||||
self._n_steps = n_steps
|
||||
self._n_etypes = n_etypes
|
||||
if not bias:
|
||||
raise KeyError("MXNet do not support disabling bias in GRUCell.")
|
||||
with self.name_scope():
|
||||
self.linears = nn.Sequential()
|
||||
for _ in range(n_etypes):
|
||||
self.linears.add(
|
||||
nn.Dense(
|
||||
out_feats,
|
||||
weight_initializer=mx.init.Xavier(),
|
||||
in_units=out_feats,
|
||||
)
|
||||
)
|
||||
self.gru = gluon.rnn.GRUCell(out_feats, input_size=out_feats)
|
||||
|
||||
def forward(self, graph, feat, etypes):
|
||||
"""Compute Gated Graph Convolution layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
The input feature of shape :math:`(N, D_{in})` where :math:`N`
|
||||
is the number of nodes of the graph and :math:`D_{in}` is the
|
||||
input feature size.
|
||||
etypes : torch.LongTensor
|
||||
The edge type tensor of shape :math:`(E,)` where :math:`E` is
|
||||
the number of edges of the graph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
|
||||
is the output feature size.
|
||||
"""
|
||||
with graph.local_scope():
|
||||
assert graph.is_homogeneous, (
|
||||
"not a homogeneous graph; convert it with to_homogeneous "
|
||||
"and pass in the edge type as argument"
|
||||
)
|
||||
zero_pad = nd.zeros(
|
||||
(feat.shape[0], self._out_feats - feat.shape[1]),
|
||||
ctx=feat.context,
|
||||
)
|
||||
feat = nd.concat(feat, zero_pad, dim=-1)
|
||||
|
||||
for _ in range(self._n_steps):
|
||||
graph.ndata["h"] = feat
|
||||
for i in range(self._n_etypes):
|
||||
eids = (etypes.asnumpy() == i).nonzero()[0]
|
||||
eids = (
|
||||
nd.from_numpy(eids, zero_copy=True)
|
||||
.as_in_context(feat.context)
|
||||
.astype(graph.idtype)
|
||||
)
|
||||
if len(eids) > 0:
|
||||
graph.apply_edges(
|
||||
lambda edges: {
|
||||
"W_e*h": self.linears[i](edges.src["h"])
|
||||
},
|
||||
eids,
|
||||
)
|
||||
graph.update_all(fn.copy_e("W_e*h", "m"), fn.sum("m", "a"))
|
||||
a = graph.ndata.pop("a")
|
||||
feat = self.gru(a, [feat])[0]
|
||||
return feat
|
||||
@@ -0,0 +1,122 @@
|
||||
"""MXNet Module for Graph Isomorphism Network layer"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import mxnet as mx
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from .... import function as fn
|
||||
from ....utils import expand_as_pair
|
||||
|
||||
|
||||
class GINConv(nn.Block):
|
||||
r"""Graph Isomorphism 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 mxnet as mx
|
||||
>>> from mxnet import gluon
|
||||
>>> from dgl.nn import GINConv
|
||||
>>>
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> lin = gluon.nn.Dense(10)
|
||||
>>> lin.initialize(ctx=mx.cpu(0))
|
||||
>>> conv = GINConv(lin, 'max')
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat)
|
||||
>>> res
|
||||
[[ 0.44832918 -0.05283341 0.20823681 0.16020004 0.37311912 -0.03372726
|
||||
-0.05716725 -0.20730163 0.14121324 0.46083626]
|
||||
[ 0.44832918 -0.05283341 0.20823681 0.16020004 0.37311912 -0.03372726
|
||||
-0.05716725 -0.20730163 0.14121324 0.46083626]
|
||||
[ 0.44832918 -0.05283341 0.20823681 0.16020004 0.37311912 -0.03372726
|
||||
-0.05716725 -0.20730163 0.14121324 0.46083626]
|
||||
[ 0.44832918 -0.05283341 0.20823681 0.16020004 0.37311912 -0.03372726
|
||||
-0.05716725 -0.20730163 0.14121324 0.46083626]
|
||||
[ 0.44832918 -0.05283341 0.20823681 0.16020004 0.37311912 -0.03372726
|
||||
-0.05716725 -0.20730163 0.14121324 0.46083626]
|
||||
[ 0.22416459 -0.0264167 0.10411841 0.08010002 0.18655956 -0.01686363
|
||||
-0.02858362 -0.10365082 0.07060662 0.23041813]]
|
||||
<NDArray 6x10 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, apply_func, aggregator_type, init_eps=0, learn_eps=False
|
||||
):
|
||||
super(GINConv, self).__init__()
|
||||
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)
|
||||
)
|
||||
|
||||
with self.name_scope():
|
||||
self.apply_func = apply_func
|
||||
self.eps = self.params.get(
|
||||
"eps",
|
||||
shape=(1,),
|
||||
grad_req="write" if learn_eps else "null",
|
||||
init=mx.init.Constant(init_eps),
|
||||
)
|
||||
|
||||
def forward(self, graph, feat):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute Graph Isomorphism Network layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray or a pair of mxnet.NDArray
|
||||
If a mxnet.NDArray 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 mxnet.NDArray 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
|
||||
-------
|
||||
mxnet.NDArray
|
||||
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.data(feat_dst.context)
|
||||
) * feat_dst + graph.dstdata["neigh"]
|
||||
if self.apply_func is not None:
|
||||
rst = self.apply_func(rst)
|
||||
return rst
|
||||
@@ -0,0 +1,265 @@
|
||||
"""Torch Module for GMM Conv"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet import nd
|
||||
from mxnet.gluon import nn
|
||||
from mxnet.gluon.contrib.nn import Identity
|
||||
|
||||
from .... import function as fn
|
||||
from ....base import DGLError
|
||||
from ....utils import expand_as_pair
|
||||
|
||||
|
||||
class GMMConv(nn.Block):
|
||||
r"""Gaussian Mixture Model Convolution layer from `Geometric Deep Learning on Graphs and
|
||||
Manifolds using Mixture Model CNNs <https://arxiv.org/abs/1611.08402>`__
|
||||
|
||||
.. math::
|
||||
u_{ij} &= f(x_i, x_j), x_j \in \mathcal{N}(i)
|
||||
|
||||
w_k(u) &= \exp\left(-\frac{1}{2}(u-\mu_k)^T \Sigma_k^{-1} (u - \mu_k)\right)
|
||||
|
||||
h_i^{l+1} &= \mathrm{aggregate}\left(\left\{\frac{1}{K}
|
||||
\sum_{k}^{K} w_k(u_{ij}), \forall j\in \mathcal{N}(i)\right\}\right)
|
||||
|
||||
where :math:`u` denotes the pseudo-coordinates between a vertex and one of its neighbor,
|
||||
computed using function :math:`f`, :math:`\Sigma_k^{-1}` and :math:`\mu_k` are
|
||||
learnable parameters representing the covariance matrix and mean vector of a Gaussian kernel.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
in_feats : int
|
||||
Number of input features; i.e., the number of dimensions of :math:`x_i`.
|
||||
out_feats : int
|
||||
Number of output features; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
|
||||
dim : int
|
||||
Dimensionality of pseudo-coordinte; i.e, the number of dimensions of :math:`u_{ij}`.
|
||||
n_kernels : int
|
||||
Number of kernels :math:`K`.
|
||||
aggregator_type : str
|
||||
Aggregator type (``sum``, ``mean``, ``max``). Default: ``sum``.
|
||||
residual : bool
|
||||
If True, use residual connection inside this layer. Default: ``False``.
|
||||
bias : bool
|
||||
If True, adds a learnable bias to the output. Default: ``True``.
|
||||
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.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import dgl
|
||||
>>> import numpy as np
|
||||
>>> import mxnet as mx
|
||||
>>> from dgl.nn import GMMConv
|
||||
>>>
|
||||
>>> # Case 1: Homogeneous graph
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> g = dgl.add_self_loop(g)
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = GMMConv(10, 2, 3, 2, 'mean')
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> pseudo = mx.nd.ones((12, 3))
|
||||
>>> res = conv(g, feat, pseudo)
|
||||
>>> res
|
||||
[[-0.05083769 -0.1567954 ]
|
||||
[-0.05083769 -0.1567954 ]
|
||||
[-0.05083769 -0.1567954 ]
|
||||
[-0.05083769 -0.1567954 ]
|
||||
[-0.05083769 -0.1567954 ]
|
||||
[-0.05083769 -0.1567954 ]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
|
||||
>>> # Case 2: Unidirectional bipartite graph
|
||||
>>> u = [0, 1, 0, 0, 1]
|
||||
>>> v = [0, 1, 2, 3, 2]
|
||||
>>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
|
||||
>>> u_fea = mx.nd.random.randn(2, 5)
|
||||
>>> v_fea = mx.nd.random.randn(4, 10)
|
||||
>>> pseudo = mx.nd.ones((5, 3))
|
||||
>>> conv = GMMConv((5, 10), 2, 3, 2, 'mean')
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, (u_fea, v_fea), pseudo)
|
||||
>>> res
|
||||
[[-0.1005067 -0.09494358]
|
||||
[-0.0023314 -0.07597432]
|
||||
[-0.05141905 -0.08545895]
|
||||
[-0.1005067 -0.09494358]]
|
||||
<NDArray 4x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_feats,
|
||||
out_feats,
|
||||
dim,
|
||||
n_kernels,
|
||||
aggregator_type="sum",
|
||||
residual=False,
|
||||
bias=True,
|
||||
allow_zero_in_degree=False,
|
||||
):
|
||||
super(GMMConv, self).__init__()
|
||||
|
||||
self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
|
||||
self._out_feats = out_feats
|
||||
self._dim = dim
|
||||
self._n_kernels = n_kernels
|
||||
self._allow_zero_in_degree = allow_zero_in_degree
|
||||
if aggregator_type == "sum":
|
||||
self._reducer = fn.sum
|
||||
elif aggregator_type == "mean":
|
||||
self._reducer = fn.mean
|
||||
elif aggregator_type == "max":
|
||||
self._reducer = fn.max
|
||||
else:
|
||||
raise KeyError(
|
||||
"Aggregator type {} not recognized.".format(aggregator_type)
|
||||
)
|
||||
|
||||
with self.name_scope():
|
||||
self.mu = self.params.get(
|
||||
"mu", shape=(n_kernels, dim), init=mx.init.Normal(0.1)
|
||||
)
|
||||
self.inv_sigma = self.params.get(
|
||||
"inv_sigma", shape=(n_kernels, dim), init=mx.init.Constant(1)
|
||||
)
|
||||
self.fc = nn.Dense(
|
||||
n_kernels * out_feats,
|
||||
in_units=self._in_src_feats,
|
||||
use_bias=False,
|
||||
weight_initializer=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
if residual:
|
||||
if self._in_dst_feats != out_feats:
|
||||
self.res_fc = nn.Dense(
|
||||
out_feats, in_units=self._in_dst_feats, use_bias=False
|
||||
)
|
||||
else:
|
||||
self.res_fc = Identity()
|
||||
else:
|
||||
self.res_fc = None
|
||||
|
||||
if bias:
|
||||
self.bias = self.params.get(
|
||||
"bias", shape=(out_feats,), init=mx.init.Zero()
|
||||
)
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
def set_allow_zero_in_degree(self, set_value):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
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 forward(self, graph, feat, pseudo):
|
||||
"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute Gaussian Mixture Model Convolution layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
If a single 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 tensors are given, the pair must contain two tensors of shape
|
||||
:math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
|
||||
pseudo : mxnet.NDArray
|
||||
The pseudo coordinate tensor of shape :math:`(E, D_{u})` where
|
||||
:math:`E` is the number of edges of the graph and :math:`D_{u}`
|
||||
is the dimensionality of pseudo coordinate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
|
||||
is the output feature size.
|
||||
|
||||
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``.
|
||||
"""
|
||||
if not self._allow_zero_in_degree:
|
||||
if graph.in_degrees().min() == 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)
|
||||
with graph.local_scope():
|
||||
graph.srcdata["h"] = self.fc(feat_src).reshape(
|
||||
-1, self._n_kernels, self._out_feats
|
||||
)
|
||||
E = graph.num_edges()
|
||||
# compute gaussian weight
|
||||
gaussian = -0.5 * (
|
||||
(
|
||||
pseudo.reshape(E, 1, self._dim)
|
||||
- self.mu.data(feat_src.context).reshape(
|
||||
1, self._n_kernels, self._dim
|
||||
)
|
||||
)
|
||||
** 2
|
||||
)
|
||||
gaussian = gaussian * (
|
||||
self.inv_sigma.data(feat_src.context).reshape(
|
||||
1, self._n_kernels, self._dim
|
||||
)
|
||||
** 2
|
||||
)
|
||||
gaussian = nd.exp(gaussian.sum(axis=-1, keepdims=True)) # (E, K, 1)
|
||||
graph.edata["w"] = gaussian
|
||||
graph.update_all(fn.u_mul_e("h", "w", "m"), self._reducer("m", "h"))
|
||||
rst = graph.dstdata["h"].sum(1)
|
||||
# residual connection
|
||||
if self.res_fc is not None:
|
||||
rst = rst + self.res_fc(feat_dst)
|
||||
# bias
|
||||
if self.bias is not None:
|
||||
rst = rst + self.bias.data(feat_dst.context)
|
||||
return rst
|
||||
@@ -0,0 +1,324 @@
|
||||
"""MXNet modules for graph convolutions(GCN)"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet import gluon
|
||||
|
||||
from .... import function as fn
|
||||
from ....base import DGLError
|
||||
from ....utils import expand_as_pair
|
||||
|
||||
|
||||
class GraphConv(gluon.Block):
|
||||
r"""Graph convolutional layer 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 mxnet as mx
|
||||
>>> from mxnet import gluon
|
||||
>>> import numpy as np
|
||||
>>> from dgl.nn import GraphConv
|
||||
|
||||
>>> # Case 1: Homogeneous graph
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> g = dgl.add_self_loop(g)
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = GraphConv(10, 2, norm='both', weight=True, bias=True)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat)
|
||||
>>> print(res)
|
||||
[[1.0209361 0.22472616]
|
||||
[1.1240715 0.24742813]
|
||||
[1.0209361 0.22472616]
|
||||
[1.2924911 0.28450024]
|
||||
[1.3568745 0.29867214]
|
||||
[0.7948386 0.17495811]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
|
||||
>>> # allow_zero_in_degree example
|
||||
>>> 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)
|
||||
[[1.0209361 0.22472616]
|
||||
[1.1240715 0.24742813]
|
||||
[1.0209361 0.22472616]
|
||||
[1.2924911 0.28450024]
|
||||
[1.3568745 0.29867214]
|
||||
[0. 0.]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
|
||||
>>> # Case 2: Unidirectional bipartite graph
|
||||
>>> u = [0, 1, 0, 0, 1]
|
||||
>>> v = [0, 1, 2, 3, 2]
|
||||
>>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
|
||||
>>> u_fea = mx.nd.random.randn(2, 5)
|
||||
>>> v_fea = mx.nd.random.randn(4, 5)
|
||||
>>> conv = GraphConv(5, 2, norm='both', weight=True, bias=True)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, (u_fea, v_fea))
|
||||
>>> res
|
||||
[[ 0.26967263 0.308129 ]
|
||||
[ 0.05143356 -0.11355402]
|
||||
[ 0.22705637 0.1375853 ]
|
||||
[ 0.26967263 0.308129 ]]
|
||||
<NDArray 4x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
with self.name_scope():
|
||||
if weight:
|
||||
self.weight = self.params.get(
|
||||
"weight",
|
||||
shape=(in_feats, out_feats),
|
||||
init=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
else:
|
||||
self.weight = None
|
||||
|
||||
if bias:
|
||||
self.bias = self.params.get(
|
||||
"bias", shape=(out_feats,), init=mx.init.Zero()
|
||||
)
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
self._activation = activation
|
||||
|
||||
def set_allow_zero_in_degree(self, set_value):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
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 forward(self, graph, feat, weight=None):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute graph convolution.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray or pair of mxnet.NDArray
|
||||
If a single 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 tensors are given, the pair must contain two tensors of shape
|
||||
:math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
|
||||
|
||||
Note that in the special case of graph convolutional networks, if a pair of
|
||||
tensors is given, the latter element will not participate in computation.
|
||||
weight : torch.Tensor, optional
|
||||
Optional external weight tensor.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
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 graph.in_degrees().min() == 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 = (
|
||||
graph.out_degrees()
|
||||
.as_in_context(feat_dst.context)
|
||||
.astype("float32")
|
||||
)
|
||||
degs = mx.nd.clip(degs, a_min=1, a_max=float("inf"))
|
||||
if self._norm == "both":
|
||||
norm = mx.nd.power(degs, -0.5)
|
||||
else:
|
||||
norm = 1.0 / degs
|
||||
shp = norm.shape + (1,) * (feat_src.ndim - 1)
|
||||
norm = norm.reshape(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.data(feat_src.context)
|
||||
|
||||
if self._in_feats > self._out_feats:
|
||||
# mult W first to reduce the feature size for aggregation.
|
||||
if weight is not None:
|
||||
feat_src = mx.nd.dot(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.pop("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.pop("h")
|
||||
if weight is not None:
|
||||
rst = mx.nd.dot(rst, weight)
|
||||
|
||||
if self._norm in ["both", "right"]:
|
||||
degs = (
|
||||
graph.in_degrees()
|
||||
.as_in_context(feat_dst.context)
|
||||
.astype("float32")
|
||||
)
|
||||
degs = mx.nd.clip(degs, a_min=1, a_max=float("inf"))
|
||||
if self._norm == "both":
|
||||
norm = mx.nd.power(degs, -0.5)
|
||||
else:
|
||||
norm = 1.0 / degs
|
||||
shp = norm.shape + (1,) * (feat_dst.ndim - 1)
|
||||
norm = norm.reshape(shp)
|
||||
rst = rst * norm
|
||||
|
||||
if self.bias is not None:
|
||||
rst = rst + self.bias.data(rst.context)
|
||||
|
||||
if self._activation is not None:
|
||||
rst = self._activation(rst)
|
||||
|
||||
return rst
|
||||
|
||||
def __repr__(self):
|
||||
summary = "GraphConv("
|
||||
summary += "in={:d}, out={:d}, normalization={}, activation={}".format(
|
||||
self._in_feats, self._out_feats, self._norm, self._activation
|
||||
)
|
||||
summary += ")"
|
||||
return summary
|
||||
@@ -0,0 +1,180 @@
|
||||
"""MXNet Module for NNConv layer"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import mxnet as mx
|
||||
from mxnet.gluon import nn
|
||||
from mxnet.gluon.contrib.nn import Identity
|
||||
|
||||
from .... import function as fn
|
||||
from ....utils import expand_as_pair
|
||||
|
||||
|
||||
class NNConv(nn.Block):
|
||||
r"""Graph Convolution layer from `Neural Message Passing
|
||||
for Quantum Chemistry <https://arxiv.org/pdf/1704.01212.pdf>`__
|
||||
|
||||
.. math::
|
||||
h_{i}^{l+1} = h_{i}^{l} + \mathrm{aggregate}\left(\left\{
|
||||
f_\Theta (e_{ij}) \cdot h_j^{l}, j\in \mathcal{N}(i) \right\}\right)
|
||||
|
||||
where :math:`e_{ij}` is the edge feature, :math:`f_\Theta` is a function
|
||||
with learnable parameters.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
in_feats : int
|
||||
Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
|
||||
NN 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 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.
|
||||
out_feats : int
|
||||
Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
|
||||
edge_func : callable activation function/layer
|
||||
Maps each edge feature to a vector of shape
|
||||
``(in_feats * out_feats)`` as weight to compute
|
||||
messages.
|
||||
Also is the :math:`f_\Theta` in the formula.
|
||||
aggregator_type : str
|
||||
Aggregator type to use (``sum``, ``mean`` or ``max``).
|
||||
residual : bool, optional
|
||||
If True, use residual connection. Default: ``False``.
|
||||
bias : bool, optional
|
||||
If True, adds a learnable bias to the output. Default: ``True``.
|
||||
|
||||
Examples
|
||||
--------
|
||||
>>> import dgl
|
||||
>>> import numpy as np
|
||||
>>> import mxnet as mx
|
||||
>>> from mxnet import gluon
|
||||
>>> from dgl.nn import NNConv
|
||||
>>>
|
||||
>>> # Case 1: Homogeneous graph
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> g = dgl.add_self_loop(g)
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> lin = gluon.nn.Dense(20)
|
||||
>>> lin.initialize(ctx=mx.cpu(0))
|
||||
>>> def edge_func(efeat):
|
||||
>>> return lin(efeat)
|
||||
>>> efeat = mx.nd.ones((12, 5))
|
||||
>>> conv = NNConv(10, 2, edge_func, 'mean')
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat, efeat)
|
||||
>>> res
|
||||
[[0.39946803 0.32098457]
|
||||
[0.39946803 0.32098457]
|
||||
[0.39946803 0.32098457]
|
||||
[0.39946803 0.32098457]
|
||||
[0.39946803 0.32098457]
|
||||
[0.39946803 0.32098457]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
|
||||
>>> # Case 2: Unidirectional bipartite graph
|
||||
>>> u = [0, 1, 0, 0, 1]
|
||||
>>> v = [0, 1, 2, 3, 2]
|
||||
>>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
|
||||
>>> u_feat = mx.nd.random.randn(2, 10)
|
||||
>>> v_feat = mx.nd.random.randn(4, 10)
|
||||
>>> conv = NNConv(10, 2, edge_func, 'mean')
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> efeat = mx.nd.ones((5, 5))
|
||||
>>> res = conv(g, (u_feat, v_feat), efeat)
|
||||
>>> res
|
||||
[[ 0.24425688 0.3238042 ]
|
||||
[-0.11651017 -0.01738572]
|
||||
[ 0.06387337 0.15320925]
|
||||
[ 0.24425688 0.3238042 ]]
|
||||
<NDArray 4x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_feats,
|
||||
out_feats,
|
||||
edge_func,
|
||||
aggregator_type,
|
||||
residual=False,
|
||||
bias=True,
|
||||
):
|
||||
super(NNConv, self).__init__()
|
||||
self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
|
||||
self._out_feats = out_feats
|
||||
if aggregator_type == "sum":
|
||||
self.reducer = fn.sum
|
||||
elif aggregator_type == "mean":
|
||||
self.reducer = fn.mean
|
||||
elif aggregator_type == "max":
|
||||
self.reducer = fn.max
|
||||
else:
|
||||
raise KeyError(
|
||||
"Aggregator type {} not recognized: ".format(aggregator_type)
|
||||
)
|
||||
self._aggre_type = aggregator_type
|
||||
|
||||
with self.name_scope():
|
||||
self.edge_nn = edge_func
|
||||
if residual:
|
||||
if self._in_dst_feats != out_feats:
|
||||
self.res_fc = nn.Dense(
|
||||
out_feats,
|
||||
in_units=self._in_dst_feats,
|
||||
use_bias=False,
|
||||
weight_initializer=mx.init.Xavier(),
|
||||
)
|
||||
else:
|
||||
self.res_fc = Identity()
|
||||
else:
|
||||
self.res_fc = None
|
||||
|
||||
if bias:
|
||||
self.bias = self.params.get(
|
||||
"bias", shape=(out_feats,), init=mx.init.Zero()
|
||||
)
|
||||
else:
|
||||
self.bias = None
|
||||
|
||||
def forward(self, graph, feat, efeat):
|
||||
r"""Compute MPNN Graph Convolution layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray or pair of mxnet.NDArray
|
||||
The input feature of shape :math:`(N, D_{in})` where :math:`N`
|
||||
is the number of nodes of the graph and :math:`D_{in}` is the
|
||||
input feature size.
|
||||
efeat : mxnet.NDArray
|
||||
The edge feature of shape :math:`(N, *)`, should fit the input
|
||||
shape requirement of ``edge_nn``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
|
||||
is the output feature size.
|
||||
"""
|
||||
with graph.local_scope():
|
||||
feat_src, feat_dst = expand_as_pair(feat, graph)
|
||||
|
||||
# (n, d_in, 1)
|
||||
graph.srcdata["h"] = feat_src.expand_dims(-1)
|
||||
# (n, d_in, d_out)
|
||||
graph.edata["w"] = self.edge_nn(efeat).reshape(
|
||||
-1, self._in_src_feats, self._out_feats
|
||||
)
|
||||
# (n, d_in, d_out)
|
||||
graph.update_all(
|
||||
fn.u_mul_e("h", "w", "m"), self.reducer("m", "neigh")
|
||||
)
|
||||
rst = graph.dstdata.pop("neigh").sum(axis=1) # (n, d_out)
|
||||
# residual connection
|
||||
if self.res_fc is not None:
|
||||
rst = rst + self.res_fc(feat_dst)
|
||||
# bias
|
||||
if self.bias is not None:
|
||||
rst = rst + self.bias.data(feat_dst.context)
|
||||
return rst
|
||||
@@ -0,0 +1,287 @@
|
||||
"""MXNet module for RelGraphConv"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
import numpy as np
|
||||
from mxnet import gluon, nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from .... import function as fn
|
||||
from .. import utils
|
||||
|
||||
|
||||
class RelGraphConv(gluon.Block):
|
||||
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 mxnet as mx
|
||||
>>> from mxnet import gluon
|
||||
>>> from dgl.nn import RelGraphConv
|
||||
>>>
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = RelGraphConv(10, 2, 3, regularizer='basis', num_bases=2)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> etype = mx.nd.array(np.array([0,1,2,0,1,2]).astype(np.int64))
|
||||
>>> res = conv(g, feat, etype)
|
||||
[[ 0.561324 0.33745846]
|
||||
[ 0.61585337 0.09992217]
|
||||
[ 0.561324 0.33745846]
|
||||
[-0.01557937 0.01227859]
|
||||
[ 0.61585337 0.09992217]
|
||||
[ 0.056508 -0.00307822]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
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
|
||||
|
||||
assert (
|
||||
low_mem is False
|
||||
), "MXNet currently does not support low-memory implementation."
|
||||
assert (
|
||||
layer_norm is False
|
||||
), "MXNet currently does not support layer norm."
|
||||
|
||||
if regularizer == "basis":
|
||||
# add basis weights
|
||||
self.weight = self.params.get(
|
||||
"weight",
|
||||
shape=(self.num_bases, self.in_feat, self.out_feat),
|
||||
init=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
if self.num_bases < self.num_rels:
|
||||
# linear combination coefficients
|
||||
self.w_comp = self.params.get(
|
||||
"w_comp",
|
||||
shape=(self.num_rels, self.num_bases),
|
||||
init=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
# 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 = self.params.get(
|
||||
"weight",
|
||||
shape=(
|
||||
self.num_rels,
|
||||
self.num_bases * self.submat_in * self.submat_out,
|
||||
),
|
||||
init=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
# 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 = self.params.get(
|
||||
"bias", shape=(out_feat,), init=mx.init.Zero()
|
||||
)
|
||||
|
||||
# weight for self loop
|
||||
if self.self_loop:
|
||||
self.loop_weight = self.params.get(
|
||||
"W_0",
|
||||
shape=(in_feat, out_feat),
|
||||
init=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
|
||||
def basis_message_func(self, edges):
|
||||
"""Message function for basis regularizer"""
|
||||
ctx = edges.src["h"].context
|
||||
if self.num_bases < self.num_rels:
|
||||
# generate all weights from bases
|
||||
weight = self.weight.data(ctx).reshape(
|
||||
self.num_bases, self.in_feat * self.out_feat
|
||||
)
|
||||
weight = nd.dot(self.w_comp.data(ctx), weight).reshape(
|
||||
self.num_rels, self.in_feat, self.out_feat
|
||||
)
|
||||
else:
|
||||
weight = self.weight.data(ctx)
|
||||
|
||||
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"""
|
||||
ctx = edges.src["h"].context
|
||||
if (
|
||||
edges.src["h"].dtype in (np.int32, np.int64)
|
||||
and len(edges.src["h"].shape) == 1
|
||||
):
|
||||
raise TypeError(
|
||||
"Block decomposition does not allow integer ID feature."
|
||||
)
|
||||
weight = self.weight.data(ctx)[edges.data["type"], :].reshape(
|
||||
-1, self.submat_in, self.submat_out
|
||||
)
|
||||
node = edges.src["h"].reshape(-1, 1, self.submat_in)
|
||||
msg = nd.batch_dot(node, weight).reshape(-1, self.out_feat)
|
||||
if "norm" in edges.data:
|
||||
msg = msg * edges.data["norm"]
|
||||
return {"msg": msg}
|
||||
|
||||
def forward(self, g, x, etypes, norm=None):
|
||||
"""
|
||||
Description
|
||||
-----------
|
||||
|
||||
Forward computation
|
||||
|
||||
Parameters
|
||||
----------
|
||||
g : DGLGraph
|
||||
The graph.
|
||||
feat : mx.ndarray.NDArray
|
||||
Input node features. Could be either
|
||||
|
||||
* :math:`(|V|, D)` dense tensor
|
||||
* :math:`(|V|,)` int64 vector, representing the categorical values of each
|
||||
node. It then treat the input feature as an one-hot encoding feature.
|
||||
etypes : mx.ndarray.NDArray
|
||||
Edge type tensor. Shape: :math:`(|E|,)`
|
||||
norm : mx.ndarray.NDArray
|
||||
Optional edge normalizer tensor. Shape: :math:`(|E|, 1)`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mx.ndarray.NDArray
|
||||
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"] = etypes
|
||||
if norm is not None:
|
||||
g.edata["norm"] = norm
|
||||
if self.self_loop:
|
||||
loop_message = utils.matmul_maybe_select(
|
||||
x, self.loop_weight.data(x.context)
|
||||
)
|
||||
|
||||
# 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.data(x.context)
|
||||
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
|
||||
@@ -0,0 +1,222 @@
|
||||
"""MXNet Module for GraphSAGE layer"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet import nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from .... import function as fn
|
||||
from ....base import DGLError
|
||||
from ....utils import check_eq_shape, expand_as_pair
|
||||
|
||||
|
||||
class SAGEConv(nn.Block):
|
||||
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 mxnet as mx
|
||||
>>> from dgl.nn import SAGEConv
|
||||
>>>
|
||||
>>> # Case 1: Homogeneous graph
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> g = dgl.add_self_loop(g)
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = SAGEConv(10, 2, 'pool')
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat)
|
||||
>>> res
|
||||
[[ 0.32144994 -0.8729614 ]
|
||||
[ 0.32144994 -0.8729614 ]
|
||||
[ 0.32144994 -0.8729614 ]
|
||||
[ 0.32144994 -0.8729614 ]
|
||||
[ 0.32144994 -0.8729614 ]
|
||||
[ 0.32144994 -0.8729614 ]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
|
||||
>>> # Case 2: Unidirectional bipartite graph
|
||||
>>> u = [0, 1, 0, 0, 1]
|
||||
>>> v = [0, 1, 2, 3, 2]
|
||||
>>> g = dgl.heterograph({('_N', '_E', '_N'):(u, v)})
|
||||
>>> u_fea = mx.nd.random.randn(2, 5)
|
||||
>>> v_fea = mx.nd.random.randn(4, 10)
|
||||
>>> conv = SAGEConv((5, 10), 2, 'pool')
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, (u_fea, v_fea))
|
||||
>>> res
|
||||
[[-0.60524774 0.7196473 ]
|
||||
[ 0.8832787 -0.5928619 ]
|
||||
[-1.8245722 1.159798 ]
|
||||
[-1.0509381 2.2239418 ]]
|
||||
<NDArray 4x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_feats,
|
||||
out_feats,
|
||||
aggregator_type="mean",
|
||||
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
|
||||
with self.name_scope():
|
||||
self.norm = norm
|
||||
self.feat_drop = nn.Dropout(feat_drop)
|
||||
self.activation = activation
|
||||
if aggregator_type == "pool":
|
||||
self.fc_pool = nn.Dense(
|
||||
self._in_src_feats,
|
||||
use_bias=bias,
|
||||
weight_initializer=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
in_units=self._in_src_feats,
|
||||
)
|
||||
if aggregator_type == "lstm":
|
||||
raise NotImplementedError
|
||||
if aggregator_type != "gcn":
|
||||
self.fc_self = nn.Dense(
|
||||
out_feats,
|
||||
use_bias=bias,
|
||||
weight_initializer=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
in_units=self._in_dst_feats,
|
||||
)
|
||||
self.fc_neigh = nn.Dense(
|
||||
out_feats,
|
||||
use_bias=bias,
|
||||
weight_initializer=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
in_units=self._in_src_feats,
|
||||
)
|
||||
|
||||
def forward(self, graph, feat):
|
||||
r"""Compute GraphSAGE layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray or pair of mxnet.NDArray
|
||||
If a single 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 tensors are given, the pair must contain two tensors of shape
|
||||
:math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
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:
|
||||
dst_neigh = mx.nd.zeros(
|
||||
(graph.number_of_dst_nodes(), self._in_src_feats)
|
||||
)
|
||||
dst_neigh = dst_neigh.as_in_context(feat_dst.context)
|
||||
graph.dstdata["neigh"] = dst_neigh
|
||||
|
||||
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 = graph.in_degrees().astype(feat_dst.dtype)
|
||||
degs = degs.as_in_context(feat_dst.context)
|
||||
h_neigh = (graph.dstdata["neigh"] + graph.dstdata["h"]) / (
|
||||
degs.expand_dims(-1) + 1
|
||||
)
|
||||
elif self._aggre_type == "pool":
|
||||
graph.srcdata["h"] = nd.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":
|
||||
raise NotImplementedError
|
||||
else:
|
||||
raise KeyError(
|
||||
"Aggregator type {} not recognized.".format(
|
||||
self._aggre_type
|
||||
)
|
||||
)
|
||||
|
||||
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
|
||||
@@ -0,0 +1,197 @@
|
||||
"""MXNet Module for Simplifying Graph Convolution layer"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet import nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from .... import function as fn
|
||||
from ....base import DGLError
|
||||
|
||||
|
||||
class SGConv(nn.Block):
|
||||
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 mxnet as mx
|
||||
>>> from dgl.nn import SGConv
|
||||
>>>
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> g = dgl.add_self_loop(g)
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = SGConv(10, 2, k=2, cached=True)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat)
|
||||
>>> res
|
||||
[[ 2.264404 -0.26684892]
|
||||
[ 2.264404 -0.26684892]
|
||||
[ 2.264404 -0.26684892]
|
||||
[ 3.2273252 -0.3803246 ]
|
||||
[ 2.247593 -0.2648679 ]
|
||||
[ 2.2644043 -0.26684904]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
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._cached = cached
|
||||
self._cached_h = None
|
||||
self._k = k
|
||||
self._allow_zero_in_degree = allow_zero_in_degree
|
||||
with self.name_scope():
|
||||
self.norm = norm
|
||||
self.fc = nn.Dense(
|
||||
out_feats,
|
||||
in_units=in_feats,
|
||||
use_bias=bias,
|
||||
weight_initializer=mx.init.Xavier(),
|
||||
)
|
||||
|
||||
def set_allow_zero_in_degree(self, set_value):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
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 forward(self, graph, feat):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute Simplifying Graph Convolution layer.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
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
|
||||
-------
|
||||
mxnet.NDArray
|
||||
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 graph.in_degrees().min() == 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 = nd.clip(
|
||||
graph.in_degrees().astype(feat.dtype), 1, float("inf")
|
||||
)
|
||||
norm = nd.power(degs, -0.5).expand_dims(1)
|
||||
norm = norm.as_in_context(feat.context)
|
||||
# compute (D^-1 A 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)
|
||||
@@ -0,0 +1,132 @@
|
||||
"""MXNet module for TAGConv"""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name
|
||||
import math
|
||||
|
||||
import mxnet as mx
|
||||
from mxnet import gluon
|
||||
|
||||
from .... import function as fn
|
||||
|
||||
|
||||
class TAGConv(gluon.Block):
|
||||
r"""Topology Adaptive Graph Convolutional layer from `Topology
|
||||
Adaptive Graph Convolutional Networks <https://arxiv.org/pdf/1710.10370.pdf>`__.
|
||||
|
||||
.. math::
|
||||
H^{K} = {\sum}_{k=0}^K (D^{-1/2} A D^{-1/2})^{k} X {\Theta}_{k},
|
||||
|
||||
where :math:`A` denotes the adjacency matrix,
|
||||
:math:`D_{ii} = \sum_{j=0} A_{ij}` its diagonal degree matrix,
|
||||
:math:`{\Theta}_{k}` denotes the linear weights to sum the results of different hops together.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
in_feats : int
|
||||
Input feature size. i.e, the number of dimensions of :math:`X`.
|
||||
out_feats : int
|
||||
Output feature size. i.e, the number of dimensions of :math:`H^{K}`.
|
||||
k: int, optional
|
||||
Number of hops :math:`K`. Default: ``2``.
|
||||
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``.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
lin : torch.Module
|
||||
The learnable linear module.
|
||||
|
||||
Example
|
||||
-------
|
||||
>>> import dgl
|
||||
>>> import numpy as np
|
||||
>>> import mxnet as mx
|
||||
>>> from mxnet import gluon
|
||||
>>> from dgl.nn import TAGConv
|
||||
>>>
|
||||
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
|
||||
>>> feat = mx.nd.ones((6, 10))
|
||||
>>> conv = TAGConv(10, 2, k=2)
|
||||
>>> conv.initialize(ctx=mx.cpu(0))
|
||||
>>> res = conv(g, feat)
|
||||
>>> res
|
||||
[[-0.86147034 0.10089529]
|
||||
[-0.86147034 0.10089529]
|
||||
[-0.86147034 0.10089529]
|
||||
[-0.9707841 0.0360311 ]
|
||||
[-0.6716844 0.02247889]
|
||||
[ 0.32964635 -0.7669234 ]]
|
||||
<NDArray 6x2 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(self, in_feats, out_feats, k=2, bias=True, activation=None):
|
||||
super(TAGConv, self).__init__()
|
||||
self.out_feats = out_feats
|
||||
self.k = k
|
||||
self.bias = bias
|
||||
self.activation = activation
|
||||
self.in_feats = in_feats
|
||||
|
||||
self.lin = self.params.get(
|
||||
"weight",
|
||||
shape=(self.in_feats * (self.k + 1), self.out_feats),
|
||||
init=mx.init.Xavier(magnitude=math.sqrt(2.0)),
|
||||
)
|
||||
if self.bias:
|
||||
self.h_bias = self.params.get(
|
||||
"bias", shape=(out_feats,), init=mx.init.Zero()
|
||||
)
|
||||
|
||||
def forward(self, graph, feat):
|
||||
r"""
|
||||
|
||||
Description
|
||||
-----------
|
||||
Compute topology adaptive graph convolution.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
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
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
|
||||
is size of output feature.
|
||||
"""
|
||||
with graph.local_scope():
|
||||
assert graph.is_homogeneous, "Graph is not homogeneous"
|
||||
|
||||
degs = graph.in_degrees().astype("float32")
|
||||
norm = mx.nd.power(
|
||||
mx.nd.clip(degs, a_min=1, a_max=float("inf")), -0.5
|
||||
)
|
||||
shp = norm.shape + (1,) * (feat.ndim - 1)
|
||||
norm = norm.reshape(shp).as_in_context(feat.context)
|
||||
|
||||
rst = feat
|
||||
for _ in range(self.k):
|
||||
rst = rst * norm
|
||||
graph.ndata["h"] = rst
|
||||
|
||||
graph.update_all(
|
||||
fn.copy_u(u="h", out="m"), fn.sum(msg="m", out="h")
|
||||
)
|
||||
rst = graph.ndata["h"]
|
||||
rst = rst * norm
|
||||
feat = mx.nd.concat(feat, rst, dim=-1)
|
||||
|
||||
rst = mx.nd.dot(feat, self.lin.data(feat.context))
|
||||
if self.bias is not None:
|
||||
rst = rst + self.h_bias.data(rst.context)
|
||||
|
||||
if self.activation is not None:
|
||||
rst = self.activation(rst)
|
||||
|
||||
return rst
|
||||
@@ -0,0 +1,327 @@
|
||||
"""MXNet modules for graph global pooling."""
|
||||
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
|
||||
from mxnet import gluon, nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
from ...readout import (
|
||||
broadcast_nodes,
|
||||
max_nodes,
|
||||
mean_nodes,
|
||||
softmax_nodes,
|
||||
sum_nodes,
|
||||
topk_nodes,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"SumPooling",
|
||||
"AvgPooling",
|
||||
"MaxPooling",
|
||||
"SortPooling",
|
||||
"GlobalAttentionPooling",
|
||||
"Set2Set",
|
||||
]
|
||||
|
||||
|
||||
class SumPooling(nn.Block):
|
||||
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 forward(self, graph, feat):
|
||||
r"""Compute sum pooling.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
The input feature with shape :math:`(N, *)` where
|
||||
:math:`N` is the number of nodes in the graph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
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")
|
||||
graph.ndata.pop("h")
|
||||
return readout
|
||||
|
||||
def __repr__(self):
|
||||
return "SumPooling()"
|
||||
|
||||
|
||||
class AvgPooling(nn.Block):
|
||||
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 forward(self, graph, feat):
|
||||
r"""Compute average pooling.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
The input feature with shape :math:`(N, *)` where
|
||||
:math:`N` is the number of nodes in the graph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
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")
|
||||
graph.ndata.pop("h")
|
||||
return readout
|
||||
|
||||
def __repr__(self):
|
||||
return "AvgPooling()"
|
||||
|
||||
|
||||
class MaxPooling(nn.Block):
|
||||
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 forward(self, graph, feat):
|
||||
r"""Compute max pooling.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
The input feature with shape :math:`(N, *)` where
|
||||
:math:`N` is the number of nodes in the graph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
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")
|
||||
graph.ndata.pop("h")
|
||||
return readout
|
||||
|
||||
def __repr__(self):
|
||||
return "MaxPooling()"
|
||||
|
||||
|
||||
class SortPooling(nn.Block):
|
||||
r"""Pooling layer 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 forward(self, graph, feat):
|
||||
r"""Compute sort pooling.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
The input node feature with shape :math:`(N, D)` where
|
||||
:math:`N` is the number of nodes in the graph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature with shape :math:`(B, k * D)`, where
|
||||
:math:`B` refers to the batch size.
|
||||
"""
|
||||
# Sort the feature of each node in ascending order.
|
||||
with graph.local_scope():
|
||||
feat = feat.sort(axis=-1)
|
||||
graph.ndata["h"] = feat
|
||||
# Sort nodes according to their last features.
|
||||
ret = topk_nodes(graph, "h", self.k, sortby=-1)[0].reshape(
|
||||
-1, self.k * feat.shape[-1]
|
||||
)
|
||||
return ret
|
||||
|
||||
def __repr__(self):
|
||||
return "SortPooling(k={})".format(self.k)
|
||||
|
||||
|
||||
class GlobalAttentionPooling(nn.Block):
|
||||
r"""Global Attention Pooling layer 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 : gluon.nn.Block
|
||||
A neural network that computes attention scores for each feature.
|
||||
feat_nn : gluon.nn.Block, 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__()
|
||||
with self.name_scope():
|
||||
self.gate_nn = gate_nn
|
||||
self.feat_nn = feat_nn
|
||||
|
||||
def forward(self, graph, feat):
|
||||
r"""Compute global attention pooling.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
The input node feature with shape :math:`(N, D)` where
|
||||
:math:`N` is the number of nodes in the graph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature with shape :math:`(B, D)`, 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["r"] = feat * gate
|
||||
readout = sum_nodes(graph, "r")
|
||||
|
||||
return readout
|
||||
|
||||
|
||||
class Set2Set(nn.Block):
|
||||
r"""Set2Set operator from `Order Matters: Sequence to sequence for sets
|
||||
<https://arxiv.org/pdf/1511.06391.pdf>`__
|
||||
|
||||
For each individual graph in the batch, set2set computes
|
||||
|
||||
.. math::
|
||||
q_t &= \mathrm{LSTM} (q^*_{t-1})
|
||||
|
||||
\alpha_{i,t} &= \mathrm{softmax}(x_i \cdot q_t)
|
||||
|
||||
r_t &= \sum_{i=1}^N \alpha_{i,t} x_i
|
||||
|
||||
q^*_t &= q_t \Vert r_t
|
||||
|
||||
for this graph.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
input_dim : int
|
||||
Size of each input sample
|
||||
n_iters : int
|
||||
Number of iterations.
|
||||
n_layers : int
|
||||
Number of recurrent layers.
|
||||
"""
|
||||
|
||||
def __init__(self, input_dim, n_iters, n_layers):
|
||||
super(Set2Set, self).__init__()
|
||||
self.input_dim = input_dim
|
||||
self.output_dim = 2 * input_dim
|
||||
self.n_iters = n_iters
|
||||
self.n_layers = n_layers
|
||||
with self.name_scope():
|
||||
self.lstm = gluon.rnn.LSTM(
|
||||
self.input_dim, num_layers=n_layers, input_size=self.output_dim
|
||||
)
|
||||
|
||||
def forward(self, graph, feat):
|
||||
r"""Compute set2set pooling.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph
|
||||
The graph.
|
||||
feat : mxnet.NDArray
|
||||
The input node feature with shape :math:`(N, D)` where
|
||||
:math:`N` is the number of nodes in the graph.
|
||||
|
||||
Returns
|
||||
-------
|
||||
mxnet.NDArray
|
||||
The output feature with shape :math:`(B, D)`, where
|
||||
:math:`B` refers to the batch size.
|
||||
"""
|
||||
with graph.local_scope():
|
||||
batch_size = graph.batch_size
|
||||
|
||||
h = (
|
||||
nd.zeros(
|
||||
(self.n_layers, batch_size, self.input_dim),
|
||||
ctx=feat.context,
|
||||
),
|
||||
nd.zeros(
|
||||
(self.n_layers, batch_size, self.input_dim),
|
||||
ctx=feat.context,
|
||||
),
|
||||
)
|
||||
q_star = nd.zeros((batch_size, self.output_dim), ctx=feat.context)
|
||||
|
||||
for _ in range(self.n_iters):
|
||||
q, h = self.lstm(q_star.expand_dims(axis=0), h)
|
||||
q = q.reshape((batch_size, self.input_dim))
|
||||
e = (feat * broadcast_nodes(graph, q)).sum(
|
||||
axis=-1, keepdims=True
|
||||
)
|
||||
graph.ndata["e"] = e
|
||||
alpha = softmax_nodes(graph, "e")
|
||||
graph.ndata["r"] = feat * alpha
|
||||
readout = sum_nodes(graph, "r")
|
||||
q_star = nd.concat(q, readout, dim=-1)
|
||||
|
||||
return q_star
|
||||
|
||||
def __repr__(self):
|
||||
summary = "Set2Set("
|
||||
summary += "in={}, out={}, " "n_iters={}, n_layers={}".format(
|
||||
self.input_dim, self.output_dim, self.n_iters, self.n_layers
|
||||
)
|
||||
summary += ")"
|
||||
return summary
|
||||
@@ -0,0 +1,254 @@
|
||||
"""Heterograph NN modules"""
|
||||
from mxnet import nd
|
||||
from mxnet.gluon import nn
|
||||
|
||||
__all__ = ["HeteroGraphConv"]
|
||||
|
||||
|
||||
class HeteroGraphConv(nn.Block):
|
||||
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 mxnet.ndarray as nd
|
||||
>>> h1 = {'user' : nd.random.randn(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'])
|
||||
|
||||
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 = mx.nd.stack(*tensors, axis=0)
|
||||
return mx.nd.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__()
|
||||
with self.name_scope():
|
||||
for name, mod in mods.items():
|
||||
self.register_child(mod, name)
|
||||
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 forward(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 __repr__(self):
|
||||
summary = "HeteroGraphConv({\n"
|
||||
for name, mod in self.mods.items():
|
||||
summary += " {} : {},\n".format(name, mod)
|
||||
summary += "\n})"
|
||||
return summary
|
||||
|
||||
|
||||
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 = nd.sum
|
||||
elif agg == "max":
|
||||
fn = nd.max
|
||||
elif agg == "min":
|
||||
fn = nd.min
|
||||
elif agg == "mean":
|
||||
fn = nd.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 nd.stack(*inputs, axis=1)
|
||||
|
||||
return stack_agg
|
||||
else:
|
||||
|
||||
def aggfn(inputs, dsttype): # pylint: disable=unused-argument
|
||||
if len(inputs) == 0:
|
||||
return None
|
||||
stacked = nd.stack(*inputs, axis=0)
|
||||
return fn(stacked, axis=0)
|
||||
|
||||
return aggfn
|
||||
@@ -0,0 +1,3 @@
|
||||
"""Gluon layer for graph related softmax."""
|
||||
# pylint: disable= unused-import
|
||||
from ..functional import edge_softmax
|
||||
@@ -0,0 +1,239 @@
|
||||
"""Utilities for pytorch NN package"""
|
||||
# pylint: disable=no-member, invalid-name
|
||||
|
||||
import numpy as np
|
||||
from mxnet import gluon, nd
|
||||
|
||||
from ... import DGLGraph
|
||||
|
||||
|
||||
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 : mxnet.NDArray
|
||||
lhs tensor
|
||||
B : mxnet.NDArray
|
||||
rhs tensor
|
||||
|
||||
Returns
|
||||
-------
|
||||
C : mxnet.NDArray
|
||||
result tensor
|
||||
"""
|
||||
if A.dtype in (np.int32, np.int64) and len(A.shape) == 1:
|
||||
return nd.take(B, A, axis=0)
|
||||
else:
|
||||
return nd.dot(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 : mxnet.NDArray
|
||||
lhs tensor
|
||||
B : mxnet.NDArray
|
||||
rhs tensor
|
||||
index : mxnet.NDArray
|
||||
index tensor
|
||||
|
||||
Returns
|
||||
-------
|
||||
C : mxnet.NDArray
|
||||
return tensor
|
||||
"""
|
||||
if A.dtype in (np.int32, np.int64) and len(A.shape) == 1:
|
||||
return B[index, A, :]
|
||||
else:
|
||||
BB = nd.take(B, index, axis=0)
|
||||
return nd.batch_dot(A.expand_dims(1), BB).squeeze(1)
|
||||
|
||||
|
||||
def normalize(x, p=2, axis=1, eps=1e-12):
|
||||
r"""Performs :math:`L_p` normalization of inputs over specified dimension.
|
||||
|
||||
For a tensor :attr:`input` of sizes :math:`(n_0, ..., n_{dim}, ..., n_k)`, each
|
||||
:math:`n_{dim}` -element vector :math:`v` along dimension :attr:`dim` is transformed as
|
||||
|
||||
.. math::
|
||||
v = \frac{v}{\max(\lVert v \rVert_p, \epsilon)}.
|
||||
|
||||
With the default arguments it uses the Euclidean norm over vectors along dimension
|
||||
:math:`1` for normalization.
|
||||
|
||||
Args:
|
||||
x: input ndarray of any shape
|
||||
ord (float): the exponent value in the norm formulation. Default: 2
|
||||
dim (int): the dimension to reduce. Default: 1
|
||||
eps (float): small value to avoid division by zero. Default: 1e-12
|
||||
"""
|
||||
denom = nd.clip(
|
||||
nd.norm(x, ord=p, axis=axis, keepdims=True), eps, float("inf")
|
||||
)
|
||||
return x / denom
|
||||
|
||||
|
||||
class Sequential(gluon.nn.Sequential):
|
||||
r"""A squential container for stacking graph neural network blocks
|
||||
|
||||
We support two modes: sequentially apply GNN blocks on the same graph or
|
||||
a list of given graphs. In the second case, the number of graphs equals the
|
||||
number of blocks inside this container.
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
Mode 1: sequentially apply GNN modules on the same graph
|
||||
|
||||
>>> import dgl
|
||||
>>> from mxnet import nd
|
||||
>>> from mxnet.gluon import nn
|
||||
>>> import dgl.function as fn
|
||||
>>> from dgl.nn.mxnet import Sequential
|
||||
>>> class ExampleLayer(nn.Block):
|
||||
>>> def __init__(self, **kwargs):
|
||||
>>> super().__init__(**kwargs)
|
||||
>>> def forward(self, graph, n_feat, e_feat):
|
||||
>>> with graph.local_scope():
|
||||
>>> graph.ndata['h'] = n_feat
|
||||
>>> graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
|
||||
>>> n_feat += graph.ndata['h']
|
||||
>>> graph.apply_edges(fn.u_add_v('h', 'h', 'e'))
|
||||
>>> e_feat += graph.edata['e']
|
||||
>>> return n_feat, e_feat
|
||||
>>>
|
||||
>>> g = dgl.DGLGraph()
|
||||
>>> g.add_nodes(3)
|
||||
>>> g.add_edges([0, 1, 2, 0, 1, 2, 0, 1, 2], [0, 0, 0, 1, 1, 1, 2, 2, 2])
|
||||
>>> net = Sequential()
|
||||
>>> net.add(ExampleLayer())
|
||||
>>> net.add(ExampleLayer())
|
||||
>>> net.add(ExampleLayer())
|
||||
>>> net.initialize()
|
||||
>>> n_feat = nd.random.randn(3, 4)
|
||||
>>> e_feat = nd.random.randn(9, 4)
|
||||
>>> net(g, n_feat, e_feat)
|
||||
(
|
||||
[[ 12.412863 99.61184 21.472883 -57.625923 ]
|
||||
[ 10.08097 100.68611 20.627377 -60.13458 ]
|
||||
[ 11.7912245 101.80654 22.427956 -58.32772 ]]
|
||||
<NDArray 3x4 @cpu(0)>,
|
||||
[[ 21.818504 198.12076 42.72387 -115.147736]
|
||||
[ 23.070837 195.49811 43.42292 -116.17203 ]
|
||||
[ 24.330334 197.10927 42.40048 -118.06538 ]
|
||||
[ 21.907919 199.11469 42.1187 -115.35658 ]
|
||||
[ 22.849625 198.79213 43.866085 -113.65381 ]
|
||||
[ 20.926125 198.116 42.64334 -114.246704]
|
||||
[ 23.003159 197.06662 41.796425 -117.14977 ]
|
||||
[ 21.391375 198.3348 41.428078 -116.30361 ]
|
||||
[ 21.291483 200.0701 40.8239 -118.07314 ]]
|
||||
<NDArray 9x4 @cpu(0)>)
|
||||
|
||||
Mode 2: sequentially apply GNN modules on different graphs
|
||||
|
||||
>>> import dgl
|
||||
>>> from mxnet import nd
|
||||
>>> from mxnet.gluon import nn
|
||||
>>> import dgl.function as fn
|
||||
>>> import networkx as nx
|
||||
>>> from dgl.nn.mxnet import Sequential
|
||||
>>> class ExampleLayer(nn.Block):
|
||||
>>> def __init__(self, **kwargs):
|
||||
>>> super().__init__(**kwargs)
|
||||
>>> def forward(self, graph, n_feat):
|
||||
>>> with graph.local_scope():
|
||||
>>> graph.ndata['h'] = n_feat
|
||||
>>> graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
|
||||
>>> n_feat += graph.ndata['h']
|
||||
>>> return n_feat.reshape(graph.num_nodes() // 2, 2, -1).sum(1)
|
||||
>>>
|
||||
>>> g1 = dgl.DGLGraph(nx.erdos_renyi_graph(32, 0.05))
|
||||
>>> g2 = dgl.DGLGraph(nx.erdos_renyi_graph(16, 0.2))
|
||||
>>> g3 = dgl.DGLGraph(nx.erdos_renyi_graph(8, 0.8))
|
||||
>>> net = Sequential()
|
||||
>>> net.add(ExampleLayer())
|
||||
>>> net.add(ExampleLayer())
|
||||
>>> net.add(ExampleLayer())
|
||||
>>> net.initialize()
|
||||
>>> n_feat = nd.random.randn(32, 4)
|
||||
>>> net([g1, g2, g3], n_feat)
|
||||
[[-101.289566 -22.584694 -89.25348 -151.6447 ]
|
||||
[-130.74239 -49.494812 -120.250854 -199.81546 ]
|
||||
[-112.32089 -50.036713 -116.13266 -190.38638 ]
|
||||
[-119.23065 -26.78553 -111.11185 -166.08322 ]]
|
||||
<NDArray 4x4 @cpu(0)>
|
||||
"""
|
||||
|
||||
def __init__(self, prefix=None, params=None):
|
||||
super(Sequential, self).__init__(prefix=prefix, params=params)
|
||||
|
||||
def forward(self, graph, *feats):
|
||||
r"""Sequentially apply modules to the input.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
graph : DGLGraph or list of DGLGraphs
|
||||
The graph(s) to apply modules on.
|
||||
|
||||
*feats :
|
||||
Input features.
|
||||
The output of :math:`i`-th block should match that of the input
|
||||
of :math:`(i+1)`-th block.
|
||||
"""
|
||||
if isinstance(graph, list):
|
||||
for graph_i, module in zip(graph, self):
|
||||
if not isinstance(feats, tuple):
|
||||
feats = (feats,)
|
||||
feats = module(graph_i, *feats)
|
||||
elif isinstance(graph, DGLGraph):
|
||||
for module in self:
|
||||
if not isinstance(feats, tuple):
|
||||
feats = (feats,)
|
||||
feats = module(graph, *feats)
|
||||
else:
|
||||
raise TypeError(
|
||||
"The first argument of forward must be a DGLGraph"
|
||||
" or a list of DGLGraph s"
|
||||
)
|
||||
return feats
|
||||
Reference in New Issue
Block a user