dmlc--dgl
Zihao Ye
af61e2fbb4
* init gat * fix * gin * 7 nn modules * rename & lint * upd * upd * fix lint * upd test * upd * lint * shape check * upd * lint * address comments * update tensorflow Co-authored-by: Quan Gan <coin2028@hotmail.com> Co-authored-by: Jinjing Zhou <VoVAllen@users.noreply.github.com> Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>
76 行
2.6 KiB
Python
76 行
2.6 KiB
Python
"""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 ..softmax import edge_softmax
|
|
from ..utils import normalize
|
|
from ....utils import expand_as_pair
|
|
|
|
|
|
class AGNNConv(nn.Block):
|
|
r"""Attention-based Graph Neural Network layer from paper `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))
|
|
|
|
Parameters
|
|
----------
|
|
init_beta : float, optional
|
|
The :math:`\beta` in the formula.
|
|
learn_beta : bool, optional
|
|
If True, :math:`\beta` will be learnable parameter.
|
|
"""
|
|
def __init__(self,
|
|
init_beta=1.,
|
|
learn_beta=True):
|
|
super(AGNNConv, self).__init__()
|
|
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 forward(self, graph, feat):
|
|
r"""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 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.
|
|
"""
|
|
graph = graph.local_var()
|
|
|
|
feat_src, feat_dst = expand_as_pair(feat)
|
|
graph.srcdata['h'] = feat_src
|
|
graph.srcdata['norm_h'] = normalize(feat_src, p=2, axis=-1)
|
|
if isinstance(feat, tuple):
|
|
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')
|