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

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dgl/nn/pytorch/__init__.py
+13
View File
@@ -0,0 +1,13 @@
"""Package for pytorch-specific NN modules."""
from .conv import *
from .explain import *
from .link import *
from .linear import *
from .glob import *
from .softmax import *
from .factory import *
from .hetero import *
from .sparse_emb import NodeEmbedding
from .utils import JumpingKnowledge, LabelPropagation, Sequential, WeightBasis
from .network_emb import *
from .gt import *
python/dgl/nn/pytorch/conv/__init__.py
+78
View File
@@ -0,0 +1,78 @@
"""Torch modules for graph convolutions."""
# pylint: disable= no-member, arguments-differ, invalid-name
from .agnnconv import AGNNConv
from .appnpconv import APPNPConv
from .atomicconv import AtomicConv
from .cfconv import CFConv
from .chebconv import ChebConv
from .cugraph_gatconv import CuGraphGATConv
from .cugraph_relgraphconv import CuGraphRelGraphConv
from .cugraph_sageconv import CuGraphSAGEConv
from .densechebconv import DenseChebConv
from .densegraphconv import DenseGraphConv
from .densesageconv import DenseSAGEConv
from .dgnconv import DGNConv
from .dotgatconv import DotGatConv
from .edgeconv import EdgeConv
from .edgegatconv import EdgeGATConv
from .egatconv import EGATConv
from .egnnconv import EGNNConv
from .gatconv import GATConv
from .gatedgcnconv import GatedGCNConv
from .gatedgraphconv import GatedGraphConv
from .gatv2conv import GATv2Conv
from .gcn2conv import GCN2Conv
from .ginconv import GINConv
from .gineconv import GINEConv
from .gmmconv import GMMConv
from .graphconv import EdgeWeightNorm, GraphConv
from .grouprevres import GroupRevRes
from .hgtconv import HGTConv
from .nnconv import NNConv
from .pnaconv import PNAConv
from .relgraphconv import RelGraphConv
from .sageconv import SAGEConv
from .sgconv import SGConv
from .tagconv import TAGConv
from .twirlsconv import TWIRLSConv, TWIRLSUnfoldingAndAttention
__all__ = [
"GraphConv",
"EdgeWeightNorm",
"GATConv",
"GATv2Conv",
"EGATConv",
"EdgeGATConv",
"TAGConv",
"RelGraphConv",
"SAGEConv",
"SGConv",
"APPNPConv",
"GINConv",
"GINEConv",
"GatedGraphConv",
"GatedGCNConv",
"GMMConv",
"ChebConv",
"AGNNConv",
"NNConv",
"DenseGraphConv",
"DenseSAGEConv",
"DenseChebConv",
"EdgeConv",
"AtomicConv",
"CFConv",
"DotGatConv",
"TWIRLSConv",
"TWIRLSUnfoldingAndAttention",
"GCN2Conv",
"HGTConv",
"GroupRevRes",
"EGNNConv",
"PNAConv",
"DGNConv",
"CuGraphGATConv",
"CuGraphRelGraphConv",
"CuGraphSAGEConv",
]
python/dgl/nn/pytorch/conv/agnnconv.py
+160
View File
@@ -0,0 +1,160 @@
"""Torch Module for Attention-based Graph Neural Network layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from torch.nn import functional as F
from .... import function as fn
from ....base import DGLError
from ....utils import expand_as_pair
from ...functional import edge_softmax
class AGNNConv(nn.Module):
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 torch as th
>>> 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 = th.ones(6, 10)
>>> conv = AGNNConv()
>>> res = conv(g, feat)
>>> res
tensor([[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.]],
grad_fn=<BinaryReduceBackward>)
"""
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
if learn_beta:
self.beta = nn.Parameter(th.Tensor([init_beta]))
else:
self.register_buffer("beta", th.Tensor([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 : torch.Tensor
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 torch.Tensor 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
-------
torch.Tensor
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() == 0).any():
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"] = F.normalize(feat_src, p=2, dim=-1)
if isinstance(feat, tuple) or graph.is_block:
graph.dstdata["norm_h"] = F.normalize(feat_dst, p=2, dim=-1)
# compute cosine distance
graph.apply_edges(fn.u_dot_v("norm_h", "norm_h", "cos"))
cos = graph.edata.pop("cos")
e = self.beta * 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")
python/dgl/nn/pytorch/conv/appnpconv.py
+123
View File
@@ -0,0 +1,123 @@
"""Torch Module for APPNPConv"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from .... import function as fn
from .graphconv import EdgeWeightNorm
class APPNPConv(nn.Module):
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 torch as th
>>> from dgl.nn import APPNPConv
>>>
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 10)
>>> conv = APPNPConv(k=3, alpha=0.5)
>>> res = conv(g, feat)
>>> print(res)
tensor([[0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536, 0.8536,
0.8536],
[0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268,
0.9268],
[0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634,
0.9634],
[0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268, 0.9268,
0.9268],
[0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634, 0.9634,
0.9634],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000,
0.5000]])
"""
def __init__(self, k, alpha, edge_drop=0.0):
super(APPNPConv, self).__init__()
self._k = k
self._alpha = alpha
self.edge_drop = nn.Dropout(edge_drop)
def forward(self, graph, feat, edge_weight=None):
r"""
Description
-----------
Compute APPNP layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor
The input feature of shape :math:`(N, *)`. :math:`N` is the
number of nodes, and :math:`*` could be of any shape.
edge_weight: torch.Tensor, optional
edge_weight to use in the message passing process. This is equivalent to
using weighted adjacency matrix in the equation above, and
:math:`\tilde{D}^{-1/2}\tilde{A} \tilde{D}^{-1/2}`
is based on :class:`dgl.nn.pytorch.conv.graphconv.EdgeWeightNorm`.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, *)` where :math:`*`
should be the same as input shape.
"""
with graph.local_scope():
if edge_weight is None:
src_norm = th.pow(
graph.out_degrees().to(feat).clamp(min=1), -0.5
)
shp = src_norm.shape + (1,) * (feat.dim() - 1)
src_norm = th.reshape(src_norm, shp).to(feat.device)
dst_norm = th.pow(
graph.in_degrees().to(feat).clamp(min=1), -0.5
)
shp = dst_norm.shape + (1,) * (feat.dim() - 1)
dst_norm = th.reshape(dst_norm, shp).to(feat.device)
else:
edge_weight = EdgeWeightNorm("both")(graph, edge_weight)
feat_0 = feat
for _ in range(self._k):
# normalization by src node
if edge_weight is None:
feat = feat * src_norm
graph.ndata["h"] = feat
w = (
th.ones(graph.num_edges(), 1)
if edge_weight is None
else edge_weight
)
graph.edata["w"] = self.edge_drop(w).to(feat.device)
graph.update_all(fn.u_mul_e("h", "w", "m"), fn.sum("m", "h"))
feat = graph.ndata.pop("h")
# normalization by dst node
if edge_weight is None:
feat = feat * dst_norm
feat = (1 - self._alpha) * feat + self._alpha * feat_0
return feat
python/dgl/nn/pytorch/conv/atomicconv.py
+301
View File
@@ -0,0 +1,301 @@
"""Torch Module for Atomic Convolution Layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import numpy as np
import torch as th
import torch.nn as nn
class RadialPooling(nn.Module):
r"""Radial pooling from `Atomic Convolutional Networks for
Predicting Protein-Ligand Binding Affinity <https://arxiv.org/abs/1703.10603>`__
We denote the distance between atom :math:`i` and :math:`j` by :math:`r_{ij}`.
A radial pooling layer transforms distances with radial filters. For radial filter
indexed by :math:`k`, it projects edge distances with
.. math::
h_{ij}^{k} = \exp(-\gamma_{k}|r_{ij}-r_{k}|^2)
If :math:`r_{ij} < c_k`,
.. math::
f_{ij}^{k} = 0.5 * \cos(\frac{\pi r_{ij}}{c_k} + 1),
else,
.. math::
f_{ij}^{k} = 0.
Finally,
.. math::
e_{ij}^{k} = h_{ij}^{k} * f_{ij}^{k}
Parameters
----------
interaction_cutoffs : float32 tensor of shape (K)
:math:`c_k` in the equations above. Roughly they can be considered as learnable cutoffs
and two atoms are considered as connected if the distance between them is smaller than
the cutoffs. K for the number of radial filters.
rbf_kernel_means : float32 tensor of shape (K)
:math:`r_k` in the equations above. K for the number of radial filters.
rbf_kernel_scaling : float32 tensor of shape (K)
:math:`\gamma_k` in the equations above. K for the number of radial filters.
"""
def __init__(
self, interaction_cutoffs, rbf_kernel_means, rbf_kernel_scaling
):
super(RadialPooling, self).__init__()
self.interaction_cutoffs = nn.Parameter(
interaction_cutoffs.reshape(-1, 1, 1), requires_grad=True
)
self.rbf_kernel_means = nn.Parameter(
rbf_kernel_means.reshape(-1, 1, 1), requires_grad=True
)
self.rbf_kernel_scaling = nn.Parameter(
rbf_kernel_scaling.reshape(-1, 1, 1), requires_grad=True
)
def forward(self, distances):
"""
Description
-----------
Apply the layer to transform edge distances.
Parameters
----------
distances : Float32 tensor of shape (E, 1)
Distance between end nodes of edges. E for the number of edges.
Returns
-------
Float32 tensor of shape (K, E, 1)
Transformed edge distances. K for the number of radial filters.
"""
scaled_euclidean_distance = (
-self.rbf_kernel_scaling * (distances - self.rbf_kernel_means) ** 2
) # (K, E, 1)
rbf_kernel_results = th.exp(scaled_euclidean_distance) # (K, E, 1)
cos_values = 0.5 * (
th.cos(np.pi * distances / self.interaction_cutoffs) + 1
) # (K, E, 1)
cutoff_values = th.where(
distances <= self.interaction_cutoffs,
cos_values,
th.zeros_like(cos_values),
) # (K, E, 1)
# Note that there appears to be an inconsistency between the paper and
# DeepChem's implementation. In the paper, the scaled_euclidean_distance first
# gets multiplied by cutoff_values, followed by exponentiation. Here we follow
# the practice of DeepChem.
return rbf_kernel_results * cutoff_values
def msg_func(edges):
"""
Description
-----------
Send messages along edges.
Parameters
----------
edges : EdgeBatch
A batch of edges.
Returns
-------
dict mapping 'm' to Float32 tensor of shape (E, K * T)
Messages computed. E for the number of edges, K for the number of
radial filters and T for the number of features to use
(types of atomic number in the paper).
"""
return {
"m": th.einsum("ij,ik->ijk", edges.src["hv"], edges.data["he"]).view(
len(edges), -1
)
}
def reduce_func(nodes):
"""
Description
-----------
Collect messages and update node representations.
Parameters
----------
nodes : NodeBatch
A batch of nodes.
Returns
-------
dict mapping 'hv_new' to Float32 tensor of shape (V, K * T)
Updated node representations. V for the number of nodes, K for the number of
radial filters and T for the number of features to use
(types of atomic number in the paper).
"""
return {"hv_new": nodes.mailbox["m"].sum(1)}
class AtomicConv(nn.Module):
r"""Atomic Convolution Layer from `Atomic Convolutional Networks for
Predicting Protein-Ligand Binding Affinity <https://arxiv.org/abs/1703.10603>`__
Denoting the type of atom :math:`i` by :math:`z_i` and the distance between atom
:math:`i` and :math:`j` by :math:`r_{ij}`.
**Distance Transformation**
An atomic convolution layer first transforms distances with radial filters and
then perform a pooling operation.
For radial filter indexed by :math:`k`, it projects edge distances with
.. math::
h_{ij}^{k} = \exp(-\gamma_{k}|r_{ij}-r_{k}|^2)
If :math:`r_{ij} < c_k`,
.. math::
f_{ij}^{k} = 0.5 * \cos(\frac{\pi r_{ij}}{c_k} + 1),
else,
.. math::
f_{ij}^{k} = 0.
Finally,
.. math::
e_{ij}^{k} = h_{ij}^{k} * f_{ij}^{k}
**Aggregation**
For each type :math:`t`, each atom collects distance information from all neighbor atoms
of type :math:`t`:
.. math::
p_{i, t}^{k} = \sum_{j\in N(i)} e_{ij}^{k} * 1(z_j == t)
Then concatenate the results for all RBF kernels and atom types.
Parameters
----------
interaction_cutoffs : float32 tensor of shape (K)
:math:`c_k` in the equations above. Roughly they can be considered as learnable cutoffs
and two atoms are considered as connected if the distance between them is smaller than
the cutoffs. K for the number of radial filters.
rbf_kernel_means : float32 tensor of shape (K)
:math:`r_k` in the equations above. K for the number of radial filters.
rbf_kernel_scaling : float32 tensor of shape (K)
:math:`\gamma_k` in the equations above. K for the number of radial filters.
features_to_use : None or float tensor of shape (T)
In the original paper, these are atomic numbers to consider, representing the types
of atoms. T for the number of types of atomic numbers. Default to None.
Note
----
* This convolution operation is designed for molecular graphs in Chemistry, but it might
be possible to extend it to more general graphs.
* There seems to be an inconsistency about the definition of :math:`e_{ij}^{k}` in the
paper and the author's implementation. We follow the author's implementation. In the
paper, :math:`e_{ij}^{k}` was defined as
:math:`\exp(-\gamma_{k}|r_{ij}-r_{k}|^2 * f_{ij}^{k})`.
* :math:`\gamma_{k}`, :math:`r_k` and :math:`c_k` are all learnable.
Example
-------
>>> import dgl
>>> import numpy as np
>>> import torch as th
>>> from dgl.nn import AtomicConv
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 1)
>>> edist = th.ones(6, 1)
>>> interaction_cutoffs = th.ones(3).float() * 2
>>> rbf_kernel_means = th.ones(3).float()
>>> rbf_kernel_scaling = th.ones(3).float()
>>> conv = AtomicConv(interaction_cutoffs, rbf_kernel_means, rbf_kernel_scaling)
>>> res = conv(g, feat, edist)
>>> res
tensor([[0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000],
[1.0000, 1.0000, 1.0000],
[0.5000, 0.5000, 0.5000],
[0.0000, 0.0000, 0.0000]], grad_fn=<ViewBackward>)
"""
def __init__(
self,
interaction_cutoffs,
rbf_kernel_means,
rbf_kernel_scaling,
features_to_use=None,
):
super(AtomicConv, self).__init__()
self.radial_pooling = RadialPooling(
interaction_cutoffs=interaction_cutoffs,
rbf_kernel_means=rbf_kernel_means,
rbf_kernel_scaling=rbf_kernel_scaling,
)
if features_to_use is None:
self.num_channels = 1
self.features_to_use = None
else:
self.num_channels = len(features_to_use)
self.features_to_use = nn.Parameter(
features_to_use, requires_grad=False
)
def forward(self, graph, feat, distances):
"""
Description
-----------
Apply the atomic convolution layer.
Parameters
----------
graph : DGLGraph
Topology based on which message passing is performed.
feat : Float32 tensor of shape :math:`(V, 1)`
Initial node features, which are atomic numbers in the paper.
:math:`V` for the number of nodes.
distances : Float32 tensor of shape :math:`(E, 1)`
Distance between end nodes of edges. E for the number of edges.
Returns
-------
Float32 tensor of shape :math:`(V, K * T)`
Updated node representations. :math:`V` for the number of nodes, :math:`K` for the
number of radial filters, and :math:`T` for the number of types of atomic numbers.
"""
with graph.local_scope():
radial_pooled_values = self.radial_pooling(distances).to(
feat
) # (K, E, 1)
if self.features_to_use is not None:
feat = (feat == self.features_to_use).to(feat) # (V, T)
graph.ndata["hv"] = feat
graph.edata["he"] = radial_pooled_values.transpose(1, 0).squeeze(
-1
) # (E, K)
graph.update_all(msg_func, reduce_func)
return graph.ndata["hv_new"].view(
graph.num_nodes(), -1
) # (V, K * T)
python/dgl/nn/pytorch/conv/cfconv.py
+145
View File
@@ -0,0 +1,145 @@
"""Torch modules for interaction blocks in SchNet"""
# pylint: disable= no-member, arguments-differ, invalid-name
import numpy as np
import torch.nn as nn
from .... import function as fn
class ShiftedSoftplus(nn.Module):
r"""Applies the element-wise function:
.. math::
\text{SSP}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x)) - \log(\text{shift})
Attributes
----------
beta : int
:math:`\beta` value for the mathematical formulation. Default to 1.
shift : int
:math:`\text{shift}` value for the mathematical formulation. Default to 2.
"""
def __init__(self, beta=1, shift=2, threshold=20):
super(ShiftedSoftplus, self).__init__()
self.shift = shift
self.softplus = nn.Softplus(beta=beta, threshold=threshold)
def forward(self, inputs):
"""
Description
-----------
Applies the activation function.
Parameters
----------
inputs : float32 tensor of shape (N, *)
* denotes any number of additional dimensions.
Returns
-------
float32 tensor of shape (N, *)
Result of applying the activation function to the input.
"""
return self.softplus(inputs) - np.log(float(self.shift))
class CFConv(nn.Module):
r"""CFConv from `SchNet: A continuous-filter convolutional neural network for
modeling quantum interactions <https://arxiv.org/abs/1706.08566>`__
It combines node and edge features in message passing and updates node representations.
.. math::
h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} h_j^{l} \circ W^{(l)}e_ij
where :math:`\circ` represents element-wise multiplication and for :math:`\text{SPP}` :
.. math::
\text{SSP}(x) = \frac{1}{\beta} * \log(1 + \exp(\beta * x)) - \log(\text{shift})
Parameters
----------
node_in_feats : int
Size for the input node features :math:`h_j^{(l)}`.
edge_in_feats : int
Size for the input edge features :math:`e_ij`.
hidden_feats : int
Size for the hidden representations.
out_feats : int
Size for the output representations :math:`h_j^{(l+1)}`.
Example
-------
>>> import dgl
>>> import numpy as np
>>> import torch as th
>>> from dgl.nn import CFConv
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> nfeat = th.ones(6, 10)
>>> efeat = th.ones(6, 5)
>>> conv = CFConv(10, 5, 3, 2)
>>> res = conv(g, nfeat, efeat)
>>> res
tensor([[-0.1209, -0.2289],
[-0.1209, -0.2289],
[-0.1209, -0.2289],
[-0.1135, -0.2338],
[-0.1209, -0.2289],
[-0.1283, -0.2240]], grad_fn=<SubBackward0>)
"""
def __init__(self, node_in_feats, edge_in_feats, hidden_feats, out_feats):
super(CFConv, self).__init__()
self.project_edge = nn.Sequential(
nn.Linear(edge_in_feats, hidden_feats),
ShiftedSoftplus(),
nn.Linear(hidden_feats, hidden_feats),
ShiftedSoftplus(),
)
self.project_node = nn.Linear(node_in_feats, hidden_feats)
self.project_out = nn.Sequential(
nn.Linear(hidden_feats, out_feats), ShiftedSoftplus()
)
def forward(self, g, node_feats, edge_feats):
"""
Description
-----------
Performs message passing and updates node representations.
Parameters
----------
g : DGLGraph
The graph.
node_feats : torch.Tensor or pair of torch.Tensor
The input node features. If a torch.Tensor is given, it represents the input
node feature of shape :math:`(N, D_{in})` where :math:`D_{in}` is size of
input feature, :math:`N` is the number of nodes.
If a pair of torch.Tensor is given, which is the case for bipartite graph,
the pair must contain two tensors of shape :math:`(N_{src}, D_{in_{src}})` and
:math:`(N_{dst}, D_{in_{dst}})` separately for the source and destination nodes.
edge_feats : torch.Tensor
The input edge feature of shape :math:`(E, edge_in_feats)`
where :math:`E` is the number of edges.
Returns
-------
torch.Tensor
The output node feature of shape :math:`(N_{out}, out_feats)`
where :math:`N_{out}` is the number of destination nodes.
"""
with g.local_scope():
if isinstance(node_feats, tuple):
node_feats_src, _ = node_feats
else:
node_feats_src = node_feats
g.srcdata["hv"] = self.project_node(node_feats_src)
g.edata["he"] = self.project_edge(edge_feats)
g.update_all(fn.u_mul_e("hv", "he", "m"), fn.sum("m", "h"))
return self.project_out(g.dstdata["h"])
python/dgl/nn/pytorch/conv/chebconv.py
+152
View File
@@ -0,0 +1,152 @@
"""Torch Module for Chebyshev Spectral Graph Convolution layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
import torch.nn.functional as F
from torch import nn
from .... import broadcast_nodes, function as fn
from ....base import dgl_warning
class ChebConv(nn.Module):
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 torch as th
>>> from dgl.nn import ChebConv
>>
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 10)
>>> conv = ChebConv(10, 2, 2)
>>> res = conv(g, feat)
>>> res
tensor([[ 0.6163, -0.1809],
[ 0.6163, -0.1809],
[ 0.6163, -0.1809],
[ 0.9698, -1.5053],
[ 0.3664, 0.7556],
[-0.2370, 3.0164]], grad_fn=<AddBackward0>)
"""
def __init__(self, in_feats, out_feats, k, activation=F.relu, bias=True):
super(ChebConv, self).__init__()
self._k = k
self._in_feats = in_feats
self._out_feats = out_feats
self.activation = activation
self.linear = nn.Linear(k * in_feats, out_feats, bias)
def forward(self, graph, feat, lambda_max=None):
r"""Compute ChebNet layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor
The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
is size of input feature, :math:`N` is the number of nodes.
lambda_max : list or tensor or None, optional.
A list(tensor) with length :math:`B`, stores the largest eigenvalue
of the normalized laplacian of each individual graph in ``graph``,
where :math:`B` is the batch size of the input graph. Default: None.
If None, this method would set the default value to 2.
One can use :func:`dgl.laplacian_lambda_max` to compute this value.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
is size of output feature.
"""
def unnLaplacian(feat, D_invsqrt, graph):
"""Operation Feat * D^-1/2 A D^-1/2"""
graph.ndata["h"] = feat * D_invsqrt
graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
return graph.ndata.pop("h") * D_invsqrt
with graph.local_scope():
D_invsqrt = th.pow(
graph.in_degrees().to(feat).clamp(min=1), -0.5
).unsqueeze(-1)
if lambda_max is None:
dgl_warning(
"lambda_max is not provided, using default value of 2. "
"Please use dgl.laplacian_lambda_max to compute the eigenvalues."
)
lambda_max = [2] * graph.batch_size
if isinstance(lambda_max, list):
lambda_max = th.Tensor(lambda_max).to(feat)
if lambda_max.dim() == 1:
lambda_max = lambda_max.unsqueeze(-1) # (B,) to (B, 1)
# broadcast from (B, 1) to (N, 1)
lambda_max = broadcast_nodes(graph, lambda_max)
re_norm = 2.0 / lambda_max
# X_0 is the raw feature, Xt is the list of X_0, X_1, ... X_t
X_0 = feat
Xt = [X_0]
# X_1(f)
if self._k > 1:
h = unnLaplacian(X_0, D_invsqrt, graph)
X_1 = -re_norm * h + X_0 * (re_norm - 1)
# Append X_1 to Xt
Xt.append(X_1)
# Xi(x), i = 2...k
for _ in range(2, self._k):
h = unnLaplacian(X_1, D_invsqrt, graph)
X_i = -2 * re_norm * h + X_1 * 2 * (re_norm - 1) - X_0
# Add X_1 to Xt
Xt.append(X_i)
X_1, X_0 = X_i, X_1
# Create the concatenation
Xt = th.cat(Xt, dim=1)
# linear projection
h = self.linear(Xt)
# activation
if self.activation:
h = self.activation(h)
return h
python/dgl/nn/pytorch/conv/cugraph_base.py
+57
View File
@@ -0,0 +1,57 @@
"""An abstract base class for cugraph-ops nn module."""
import torch
from torch import nn
class CuGraphBaseConv(nn.Module):
r"""An abstract base class for cugraph-ops nn module."""
def __init__(self):
super().__init__()
self._cached_offsets_fg = None
def reset_parameters(self):
r"""Resets all learnable parameters of the module."""
raise NotImplementedError
def forward(self, *args):
r"""Runs the forward pass of the module."""
raise NotImplementedError
def pad_offsets(self, offsets: torch.Tensor, size: int) -> torch.Tensor:
r"""Pad zero-in-degree nodes to the end of offsets to reach size.
cugraph-ops often provides two variants of aggregation functions for a
specific model: one intended for sampled-graph use cases, one for
full-graph ones. The former is in general more performant, however, it
only works when the sample size (the max of in-degrees) is small (<200),
due to the limit of GPU shared memory. For graphs with a larger max
in-degree, we need to fall back to the full-graph option, which requires
to convert a DGL block to a full graph. With the csc-representation,
this is equivalent to pad zero-in-degree nodes to the end of the offsets
array (also called indptr or colptr).
Parameters
----------
offsets :
The (monotonically increasing) index pointer array in a CSC-format
graph.
size : int
The length of offsets after padding.
Returns
-------
torch.Tensor
The augmented offsets array.
"""
if self._cached_offsets_fg is None:
self._cached_offsets_fg = torch.empty(
size, dtype=offsets.dtype, device=offsets.device
)
elif self._cached_offsets_fg.numel() < size:
self._cached_offsets_fg.resize_(size)
self._cached_offsets_fg[: offsets.numel()] = offsets
self._cached_offsets_fg[offsets.numel() : size] = offsets[-1]
return self._cached_offsets_fg[:size]
python/dgl/nn/pytorch/conv/cugraph_gatconv.py
+213
View File
@@ -0,0 +1,213 @@
"""Torch Module for graph attention network layer using the aggregation
primitives in cugraph-ops"""
# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
import torch
from torch import nn
from .cugraph_base import CuGraphBaseConv
try:
from pylibcugraphops.pytorch import SampledCSC, StaticCSC
from pylibcugraphops.pytorch.operators import mha_gat_n2n as GATConvAgg
HAS_PYLIBCUGRAPHOPS = True
except ImportError:
HAS_PYLIBCUGRAPHOPS = False
class CuGraphGATConv(CuGraphBaseConv):
r"""Graph attention layer from `Graph Attention Networks
<https://arxiv.org/pdf/1710.10903.pdf>`__, with the sparse aggregation
accelerated by cugraph-ops.
See :class:`dgl.nn.pytorch.conv.GATConv` for mathematical model.
This module depends on :code:`pylibcugraphops` package, which can be
installed via :code:`conda install -c nvidia pylibcugraphops=23.04`.
:code:`pylibcugraphops` 23.04 requires python 3.8.x or 3.10.x.
.. note::
This is an **experimental** feature.
Parameters
----------
in_feats : int
Input feature size.
out_feats : int
Output feature size.
num_heads : int
Number of heads in Multi-Head Attention.
feat_drop : float, optional
Dropout rate on feature. 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``.
bias : bool, optional
If True, learns a bias term. Defaults: ``True``.
Examples
--------
>>> import dgl
>>> import torch
>>> from dgl.nn import CuGraphGATConv
>>> device = 'cuda'
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])).to(device)
>>> g = dgl.add_self_loop(g)
>>> feat = torch.ones(6, 10).to(device)
>>> conv = CuGraphGATConv(10, 2, num_heads=3).to(device)
>>> res = conv(g, feat)
>>> res
tensor([[[ 0.2340, 1.9226],
[ 1.6477, -1.9986],
[ 1.1138, -1.9302]],
[[ 0.2340, 1.9226],
[ 1.6477, -1.9986],
[ 1.1138, -1.9302]],
[[ 0.2340, 1.9226],
[ 1.6477, -1.9986],
[ 1.1138, -1.9302]],
[[ 0.2340, 1.9226],
[ 1.6477, -1.9986],
[ 1.1138, -1.9302]],
[[ 0.2340, 1.9226],
[ 1.6477, -1.9986],
[ 1.1138, -1.9302]],
[[ 0.2340, 1.9226],
[ 1.6477, -1.9986],
[ 1.1138, -1.9302]]], device='cuda:0', grad_fn=<ViewBackward0>)
"""
MAX_IN_DEGREE_MFG = 200
def __init__(
self,
in_feats,
out_feats,
num_heads,
feat_drop=0.0,
negative_slope=0.2,
residual=False,
activation=None,
bias=True,
):
if HAS_PYLIBCUGRAPHOPS is False:
raise ModuleNotFoundError(
f"{self.__class__.__name__} requires pylibcugraphops=23.04. "
f"Install via `conda install -c nvidia 'pylibcugraphops=23.04'`."
f"pylibcugraphops requires Python 3.8 or 3.10."
)
super().__init__()
self.in_feats = in_feats
self.out_feats = out_feats
self.num_heads = num_heads
self.feat_drop = nn.Dropout(feat_drop)
self.negative_slope = negative_slope
self.activation = activation
self.fc = nn.Linear(in_feats, out_feats * num_heads, bias=False)
self.attn_weights = nn.Parameter(
torch.Tensor(2 * num_heads * out_feats)
)
if bias:
self.bias = nn.Parameter(torch.Tensor(num_heads * out_feats))
else:
self.register_buffer("bias", None)
if residual:
if in_feats == out_feats * num_heads:
self.res_fc = nn.Identity()
else:
self.res_fc = nn.Linear(
in_feats, out_feats * num_heads, bias=False
)
else:
self.register_buffer("res_fc", None)
self.reset_parameters()
def reset_parameters(self):
r"""Reinitialize learnable parameters."""
gain = nn.init.calculate_gain("relu")
nn.init.xavier_normal_(self.fc.weight, gain=gain)
nn.init.xavier_normal_(
self.attn_weights.view(2, self.num_heads, self.out_feats), gain=gain
)
if self.bias is not None:
nn.init.zeros_(self.bias)
if isinstance(self.res_fc, nn.Linear):
self.res_fc.reset_parameters()
def forward(self, g, feat, max_in_degree=None):
r"""Forward computation.
Parameters
----------
g : DGLGraph
The graph.
feat : torch.Tensor
Input features of shape :math:`(N, D_{in})`.
max_in_degree : int
Maximum in-degree of destination nodes. It is only effective when
:attr:`g` is a :class:`DGLBlock`, i.e., bipartite graph. When
:attr:`g` is generated from a neighbor sampler, the value should be
set to the corresponding :attr:`fanout`. If not given,
:attr:`max_in_degree` will be calculated on-the-fly.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, H, D_{out})` where
:math:`H` is the number of heads, and :math:`D_{out}` is size of
output feature.
"""
offsets, indices, _ = g.adj_tensors("csc")
if g.is_block:
if max_in_degree is None:
max_in_degree = g.in_degrees().max().item()
if max_in_degree < self.MAX_IN_DEGREE_MFG:
_graph = SampledCSC(
offsets,
indices,
max_in_degree,
g.num_src_nodes(),
)
else:
offsets_fg = self.pad_offsets(offsets, g.num_src_nodes() + 1)
_graph = StaticCSC(offsets_fg, indices)
else:
_graph = StaticCSC(offsets, indices)
feat = self.feat_drop(feat)
feat_transformed = self.fc(feat)
out = GATConvAgg(
feat_transformed,
self.attn_weights,
_graph,
self.num_heads,
"LeakyReLU",
self.negative_slope,
concat_heads=True,
)[: g.num_dst_nodes()].view(-1, self.num_heads, self.out_feats)
feat_dst = feat[: g.num_dst_nodes()]
if self.res_fc is not None:
out = out + self.res_fc(feat_dst).view(
-1, self.num_heads, self.out_feats
)
if self.bias is not None:
out = out + self.bias.view(-1, self.num_heads, self.out_feats)
if self.activation is not None:
out = self.activation(out)
return out
python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
+228
View File
@@ -0,0 +1,228 @@
"""Torch Module for Relational graph convolution layer using the aggregation
primitives in cugraph-ops"""
# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
import math
import torch
from torch import nn
from .cugraph_base import CuGraphBaseConv
try:
from pylibcugraphops.pytorch import HeteroCSC
from pylibcugraphops.pytorch.operators import (
agg_hg_basis_n2n_post as RelGraphConvAgg,
)
HAS_PYLIBCUGRAPHOPS = True
except ImportError:
HAS_PYLIBCUGRAPHOPS = False
class CuGraphRelGraphConv(CuGraphBaseConv):
r"""An accelerated relational graph convolution layer from `Modeling
Relational Data with Graph Convolutional Networks
<https://arxiv.org/abs/1703.06103>`__ that leverages the highly-optimized
aggregation primitives in cugraph-ops.
See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
This module depends on :code:`pylibcugraphops` package, which can be
installed via :code:`conda install -c nvidia pylibcugraphops=23.04`.
:code:`pylibcugraphops` 23.04 requires python 3.8.x or 3.10.x.
.. note::
This is an **experimental** feature.
Parameters
----------
in_feat : int
Input feature size.
out_feat : int
Output feature size.
num_rels : int
Number of relations.
regularizer : str, optional
Which weight regularizer to use ("basis" or ``None``):
- "basis" is for basis-decomposition.
- ``None`` applies no regularization.
Default: ``None``.
num_bases : int, optional
Number of bases. It comes into effect when a regularizer is applied.
Default: ``None``.
bias : bool, optional
True if bias is added. Default: ``True``.
self_loop : bool, optional
True to include self loop message. Default: ``True``.
dropout : float, optional
Dropout rate. Default: ``0.0``.
apply_norm : bool, optional
True to normalize aggregation output by the in-degree of the destination
node per edge type, i.e. :math:`|\mathcal{N}^r_i|`. Default: ``True``.
Examples
--------
>>> import dgl
>>> import torch
>>> from dgl.nn import CuGraphRelGraphConv
...
>>> device = 'cuda'
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])).to(device)
>>> feat = torch.ones(6, 10).to(device)
>>> conv = CuGraphRelGraphConv(
... 10, 2, 3, regularizer='basis', num_bases=2).to(device)
>>> etype = torch.tensor([0,1,2,0,1,2]).to(device)
>>> res = conv(g, feat, etype)
>>> res
tensor([[-1.7774, -2.0184],
[-1.4335, -2.3758],
[-1.7774, -2.0184],
[-0.4698, -3.0876],
[-1.4335, -2.3758],
[-1.4331, -2.3295]], device='cuda:0', grad_fn=<AddBackward0>)
"""
MAX_IN_DEGREE_MFG = 500
def __init__(
self,
in_feat,
out_feat,
num_rels,
regularizer=None,
num_bases=None,
bias=True,
self_loop=True,
dropout=0.0,
apply_norm=False,
):
if HAS_PYLIBCUGRAPHOPS is False:
raise ModuleNotFoundError(
f"{self.__class__.__name__} requires pylibcugraphops=23.04. "
f"Install via `conda install -c nvidia 'pylibcugraphops=23.04'`."
f"pylibcugraphops requires Python 3.8 or 3.10."
)
super().__init__()
self.in_feat = in_feat
self.out_feat = out_feat
self.num_rels = num_rels
self.apply_norm = apply_norm
self.dropout = nn.Dropout(dropout)
dim_self_loop = 1 if self_loop else 0
self.self_loop = self_loop
if regularizer is None:
self.W = nn.Parameter(
torch.Tensor(num_rels + dim_self_loop, in_feat, out_feat)
)
self.coeff = None
elif regularizer == "basis":
if num_bases is None:
raise ValueError(
'Missing "num_bases" for basis regularization.'
)
self.W = nn.Parameter(
torch.Tensor(num_bases + dim_self_loop, in_feat, out_feat)
)
self.coeff = nn.Parameter(torch.Tensor(num_rels, num_bases))
self.num_bases = num_bases
else:
raise ValueError(
f"Supported regularizer options: 'basis' or None, but got "
f"'{regularizer}'."
)
self.regularizer = regularizer
if bias:
self.bias = nn.Parameter(torch.Tensor(out_feat))
else:
self.register_parameter("bias", None)
self.reset_parameters()
def reset_parameters(self):
r"""Reinitialize learnable parameters."""
bound = 1 / math.sqrt(self.in_feat)
end = -1 if self.self_loop else None
nn.init.uniform_(self.W[:end], -bound, bound)
if self.regularizer == "basis":
nn.init.xavier_uniform_(
self.coeff, gain=nn.init.calculate_gain("relu")
)
if self.self_loop:
nn.init.xavier_uniform_(self.W[-1], nn.init.calculate_gain("relu"))
if self.bias is not None:
nn.init.zeros_(self.bias)
def forward(self, g, feat, etypes, max_in_degree=None):
r"""Forward computation.
Parameters
----------
g : DGLGraph
The graph.
feat : torch.Tensor
A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
etypes : torch.Tensor
A 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
Note that cugraph-ops only accepts edge type tensors in int32,
so any input of other integer types will be casted into int32,
thus introducing some overhead. Pass in int32 tensors directly
for best performance.
max_in_degree : int, optional
Maximum in-degree of destination nodes. It is only effective when
:attr:`g` is a :class:`DGLBlock`, i.e., bipartite graph. When
:attr:`g` is generated from a neighbor sampler, the value should be
set to the corresponding :attr:`fanout`. If not given,
:attr:`max_in_degree` will be calculated on-the-fly.
Returns
-------
torch.Tensor
New node features. Shape: :math:`(|V|, D_{out})`.
"""
offsets, indices, edge_ids = g.adj_tensors("csc")
edge_types_perm = etypes[edge_ids.long()].int()
if g.is_block:
if max_in_degree is None:
max_in_degree = g.in_degrees().max().item()
if max_in_degree < self.MAX_IN_DEGREE_MFG:
_graph = HeteroCSC(
offsets,
indices,
edge_types_perm,
g.num_src_nodes(),
self.num_rels,
)
else:
offsets_fg = self.pad_offsets(offsets, g.num_src_nodes() + 1)
_graph = HeteroCSC(
offsets_fg,
indices,
edge_types_perm,
g.num_src_nodes(),
self.num_rels,
)
else:
_graph = HeteroCSC(
offsets,
indices,
edge_types_perm,
g.num_src_nodes(),
self.num_rels,
)
h = RelGraphConvAgg(
feat,
self.coeff,
_graph,
concat_own=self.self_loop,
norm_by_out_degree=self.apply_norm,
)[: g.num_dst_nodes()]
h = h @ self.W.view(-1, self.out_feat)
if self.bias is not None:
h = h + self.bias
h = self.dropout(h)
return h
python/dgl/nn/pytorch/conv/cugraph_sageconv.py
+148
View File
@@ -0,0 +1,148 @@
"""Torch Module for GraphSAGE layer using the aggregation primitives in
cugraph-ops"""
# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
from torch import nn
from .cugraph_base import CuGraphBaseConv
try:
from pylibcugraphops.pytorch import SampledCSC, StaticCSC
from pylibcugraphops.pytorch.operators import agg_concat_n2n as SAGEConvAgg
HAS_PYLIBCUGRAPHOPS = True
except ImportError:
HAS_PYLIBCUGRAPHOPS = False
class CuGraphSAGEConv(CuGraphBaseConv):
r"""An accelerated GraphSAGE layer from `Inductive Representation Learning
on Large Graphs <https://arxiv.org/pdf/1706.02216.pdf>`__ that leverages the
highly-optimized aggregation primitives in cugraph-ops:
.. math::
h_{\mathcal{N}(i)}^{(l+1)} &= \mathrm{aggregate}
\left(\{h_{j}^{l}, \forall j \in \mathcal{N}(i) \}\right)
h_{i}^{(l+1)} &= W \cdot \mathrm{concat}
(h_{i}^{l}, h_{\mathcal{N}(i)}^{(l+1)})
This module depends on :code:`pylibcugraphops` package, which can be
installed via :code:`conda install -c nvidia pylibcugraphops=23.04`.
:code:`pylibcugraphops` 23.04 requires python 3.8.x or 3.10.x.
.. note::
This is an **experimental** feature.
Parameters
----------
in_feats : int
Input feature size.
out_feats : int
Output feature size.
aggregator_type : str
Aggregator type to use (``mean``, ``sum``, ``min``, ``max``).
feat_drop : float
Dropout rate on features, default: ``0``.
bias : bool
If True, adds a learnable bias to the output. Default: ``True``.
Examples
--------
>>> import dgl
>>> import torch
>>> from dgl.nn import CuGraphSAGEConv
>>> device = 'cuda'
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3])).to(device)
>>> g = dgl.add_self_loop(g)
>>> feat = torch.ones(6, 10).to(device)
>>> conv = CuGraphSAGEConv(10, 2, 'mean').to(device)
>>> res = conv(g, feat)
>>> res
tensor([[-1.1690, 0.1952],
[-1.1690, 0.1952],
[-1.1690, 0.1952],
[-1.1690, 0.1952],
[-1.1690, 0.1952],
[-1.1690, 0.1952]], device='cuda:0', grad_fn=<AddmmBackward0>)
"""
MAX_IN_DEGREE_MFG = 500
def __init__(
self,
in_feats,
out_feats,
aggregator_type="mean",
feat_drop=0.0,
bias=True,
):
if HAS_PYLIBCUGRAPHOPS is False:
raise ModuleNotFoundError(
f"{self.__class__.__name__} requires pylibcugraphops=23.04. "
f"Install via `conda install -c nvidia 'pylibcugraphops=23.04'`."
f"pylibcugraphops requires Python 3.8 or 3.10."
)
valid_aggr_types = {"max", "min", "mean", "sum"}
if aggregator_type not in valid_aggr_types:
raise ValueError(
f"Invalid aggregator_type. Must be one of {valid_aggr_types}. "
f"But got '{aggregator_type}' instead."
)
super().__init__()
self.in_feats = in_feats
self.out_feats = out_feats
self.aggr = aggregator_type
self.feat_drop = nn.Dropout(feat_drop)
self.linear = nn.Linear(2 * in_feats, out_feats, bias=bias)
def reset_parameters(self):
r"""Reinitialize learnable parameters."""
self.linear.reset_parameters()
def forward(self, g, feat, max_in_degree=None):
r"""Forward computation.
Parameters
----------
g : DGLGraph
The graph.
feat : torch.Tensor
Node features. Shape: :math:`(N, D_{in})`.
max_in_degree : int
Maximum in-degree of destination nodes. It is only effective when
:attr:`g` is a :class:`DGLBlock`, i.e., bipartite graph. When
:attr:`g` is generated from a neighbor sampler, the value should be
set to the corresponding :attr:`fanout`. If not given,
:attr:`max_in_degree` will be calculated on-the-fly.
Returns
-------
torch.Tensor
Output node features. Shape: :math:`(N, D_{out})`.
"""
offsets, indices, _ = g.adj_tensors("csc")
if g.is_block:
if max_in_degree is None:
max_in_degree = g.in_degrees().max().item()
if max_in_degree < self.MAX_IN_DEGREE_MFG:
_graph = SampledCSC(
offsets,
indices,
max_in_degree,
g.num_src_nodes(),
)
else:
offsets_fg = self.pad_offsets(offsets, g.num_src_nodes() + 1)
_graph = StaticCSC(offsets_fg, indices)
else:
_graph = StaticCSC(offsets, indices)
feat = self.feat_drop(feat)
h = SAGEConvAgg(feat, _graph, self.aggr)[: g.num_dst_nodes()]
h = self.linear(h)
return h
python/dgl/nn/pytorch/conv/densechebconv.py
+125
View File
@@ -0,0 +1,125 @@
"""Torch Module for DenseChebConv"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from torch.nn import init
class DenseChebConv(nn.Module):
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``.
Example
-------
>>> import dgl
>>> import numpy as np
>>> import torch as th
>>> from dgl.nn import DenseChebConv
>>>
>>> feat = th.ones(6, 10)
>>> adj = th.tensor([[0., 0., 1., 0., 0., 0.],
... [1., 0., 0., 0., 0., 0.],
... [0., 1., 0., 0., 0., 0.],
... [0., 0., 1., 0., 0., 1.],
... [0., 0., 0., 1., 0., 0.],
... [0., 0., 0., 0., 0., 0.]])
>>> conv = DenseChebConv(10, 2, 2)
>>> res = conv(adj, feat)
>>> res
tensor([[-3.3516, -2.4797],
[-3.3516, -2.4797],
[-3.3516, -2.4797],
[-4.5192, -3.0835],
[-2.5259, -2.0527],
[-0.5327, -1.0219]], grad_fn=<AddBackward0>)
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
self.W = nn.Parameter(th.Tensor(k, in_feats, out_feats))
if bias:
self.bias = nn.Parameter(th.Tensor(out_feats))
else:
self.register_buffer("bias", None)
self.reset_parameters()
def reset_parameters(self):
"""Reinitialize learnable parameters."""
if self.bias is not None:
init.zeros_(self.bias)
for i in range(self._k):
init.xavier_normal_(self.W[i], init.calculate_gain("relu"))
def forward(self, adj, feat, lambda_max=None):
r"""Compute (Dense) Chebyshev Spectral Graph Convolution layer
Parameters
----------
adj : torch.Tensor
The adjacency matrix of the graph to apply Graph Convolution on,
should be of shape :math:`(N, N)`, where a row represents the destination
and a column represents the source.
feat : torch.Tensor
The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
is size of input feature, :math:`N` is the number of nodes.
lambda_max : float or None, optional
A float value indicates the largest eigenvalue of given graph.
Default: None.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
is size of output feature.
"""
A = adj.to(feat)
num_nodes = A.shape[0]
in_degree = 1 / A.sum(dim=1).clamp(min=1).sqrt()
D_invsqrt = th.diag(in_degree)
I = th.eye(num_nodes).to(A)
L = I - D_invsqrt @ A @ D_invsqrt
if lambda_max is None:
lambda_ = th.eig(L)[0][:, 0]
lambda_max = lambda_.max()
L_hat = 2 * L / lambda_max - I
Z = [th.eye(num_nodes).to(A)]
for i in range(1, self._k):
if i == 1:
Z.append(L_hat)
else:
Z.append(2 * L_hat @ Z[-1] - Z[-2])
Zs = th.stack(Z, 0) # (k, n, n)
Zh = Zs @ feat.unsqueeze(0) @ self.W
Zh = Zh.sum(0)
if self.bias is not None:
Zh = Zh + self.bias
return Zh
python/dgl/nn/pytorch/conv/densegraphconv.py
+145
View File
@@ -0,0 +1,145 @@
"""Torch Module for DenseGraphConv"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from torch.nn import init
class DenseGraphConv(nn.Module):
"""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.
Example
-------
>>> import dgl
>>> import numpy as np
>>> import torch as th
>>> from dgl.nn import DenseGraphConv
>>>
>>> feat = th.ones(6, 10)
>>> adj = th.tensor([[0., 0., 1., 0., 0., 0.],
... [1., 0., 0., 0., 0., 0.],
... [0., 1., 0., 0., 0., 0.],
... [0., 0., 1., 0., 0., 1.],
... [0., 0., 0., 1., 0., 0.],
... [0., 0., 0., 0., 0., 0.]])
>>> conv = DenseGraphConv(10, 2)
>>> res = conv(adj, feat)
>>> res
tensor([[0.2159, 1.9027],
[0.3053, 2.6908],
[0.3053, 2.6908],
[0.3685, 3.2481],
[0.3053, 2.6908],
[0.0000, 0.0000]], grad_fn=<AddBackward0>)
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
self.weight = nn.Parameter(th.Tensor(in_feats, out_feats))
if bias:
self.bias = nn.Parameter(th.Tensor(out_feats))
else:
self.register_buffer("bias", None)
self.reset_parameters()
self._activation = activation
def reset_parameters(self):
"""Reinitialize learnable parameters."""
init.xavier_uniform_(self.weight)
if self.bias is not None:
init.zeros_(self.bias)
def forward(self, adj, feat):
r"""Compute (Dense) Graph Convolution layer.
Parameters
----------
adj : torch.Tensor
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 : torch.Tensor
The input feature.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
is size of output feature.
"""
adj = adj.to(feat)
src_degrees = adj.sum(dim=0).clamp(min=1)
dst_degrees = adj.sum(dim=1).clamp(min=1)
feat_src = feat
if self._norm == "both":
norm_src = th.pow(src_degrees, -0.5)
shp = norm_src.shape + (1,) * (feat.dim() - 1)
norm_src = th.reshape(norm_src, shp).to(feat.device)
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 = th.matmul(feat_src, self.weight)
rst = adj @ feat_src
else:
# aggregate first then mult W
rst = adj @ feat_src
rst = th.matmul(rst, self.weight)
if self._norm != "none":
if self._norm == "both":
norm_dst = th.pow(dst_degrees, -0.5)
else: # right
norm_dst = 1.0 / dst_degrees
shp = norm_dst.shape + (1,) * (feat.dim() - 1)
norm_dst = th.reshape(norm_dst, shp).to(feat.device)
rst = rst * norm_dst
if self.bias is not None:
rst = rst + self.bias
if self._activation is not None:
rst = self._activation(rst)
return rst
python/dgl/nn/pytorch/conv/densesageconv.py
+138
View File
@@ -0,0 +1,138 @@
"""Torch Module for DenseSAGEConv"""
# pylint: disable= no-member, arguments-differ, invalid-name
from torch import nn
from ....utils import check_eq_shape
class DenseSAGEConv(nn.Module):
"""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``.
Example
-------
>>> import dgl
>>> import numpy as np
>>> import torch as th
>>> from dgl.nn import DenseSAGEConv
>>>
>>> feat = th.ones(6, 10)
>>> adj = th.tensor([[0., 0., 1., 0., 0., 0.],
... [1., 0., 0., 0., 0., 0.],
... [0., 1., 0., 0., 0., 0.],
... [0., 0., 1., 0., 0., 1.],
... [0., 0., 0., 1., 0., 0.],
... [0., 0., 0., 0., 0., 0.]])
>>> conv = DenseSAGEConv(10, 2)
>>> res = conv(adj, feat)
>>> res
tensor([[1.0401, 2.1008],
[1.0401, 2.1008],
[1.0401, 2.1008],
[1.0401, 2.1008],
[1.0401, 2.1008],
[1.0401, 2.1008]], grad_fn=<AddmmBackward>)
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
self.feat_drop = nn.Dropout(feat_drop)
self.activation = activation
self.fc = nn.Linear(in_feats, out_feats, bias=bias)
self.reset_parameters()
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
Notes
-----
The linear weights :math:`W^{(l)}` are initialized using Glorot uniform initialization.
"""
gain = nn.init.calculate_gain("relu")
nn.init.xavier_uniform_(self.fc.weight, gain=gain)
def forward(self, adj, feat):
r"""
Description
-----------
Compute (Dense) Graph SAGE layer.
Parameters
----------
adj : torch.Tensor
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 : torch.Tensor or a pair of torch.Tensor
If a torch.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 torch.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, D_{in})` and :math:`(N_{out}, D_{in})`.
Returns
-------
torch.Tensor
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.to(feat_src)
in_degrees = adj.sum(dim=1, keepdim=True)
h_neigh = (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
python/dgl/nn/pytorch/conv/dgnconv.py
+265
View File
@@ -0,0 +1,265 @@
"""Torch Module for Directional Graph Networks Convolution Layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
from functools import partial
import torch
import torch.nn as nn
from .pnaconv import AGGREGATORS, PNAConv, PNAConvTower, SCALERS
def aggregate_dir_av(h, eig_s, eig_d, eig_idx):
"""directional average aggregation"""
h_mod = torch.mul(
h,
(
torch.abs(eig_s[:, :, eig_idx] - eig_d[:, :, eig_idx])
/ (
torch.sum(
torch.abs(eig_s[:, :, eig_idx] - eig_d[:, :, eig_idx]),
keepdim=True,
dim=1,
)
+ 1e-30
)
).unsqueeze(-1),
)
return torch.sum(h_mod, dim=1)
def aggregate_dir_dx(h, eig_s, eig_d, h_in, eig_idx):
"""directional derivative aggregation"""
eig_w = (
(eig_s[:, :, eig_idx] - eig_d[:, :, eig_idx])
/ (
torch.sum(
torch.abs(eig_s[:, :, eig_idx] - eig_d[:, :, eig_idx]),
keepdim=True,
dim=1,
)
+ 1e-30
)
).unsqueeze(-1)
h_mod = torch.mul(h, eig_w)
return torch.abs(torch.sum(h_mod, dim=1) - torch.sum(eig_w, dim=1) * h_in)
for k in range(1, 4):
AGGREGATORS[f"dir{k}-av"] = partial(aggregate_dir_av, eig_idx=k - 1)
AGGREGATORS[f"dir{k}-dx"] = partial(aggregate_dir_dx, eig_idx=k - 1)
class DGNConvTower(PNAConvTower):
"""A single DGN tower with modified reduce function"""
def message(self, edges):
"""message function for DGN layer"""
if self.edge_feat_size > 0:
f = torch.cat(
[edges.src["h"], edges.dst["h"], edges.data["a"]], dim=-1
)
else:
f = torch.cat([edges.src["h"], edges.dst["h"]], dim=-1)
return {
"msg": self.M(f),
"eig_s": edges.src["eig"],
"eig_d": edges.dst["eig"],
}
def reduce_func(self, nodes):
"""reduce function for DGN layer"""
h_in = nodes.data["h"]
eig_s = nodes.mailbox["eig_s"]
eig_d = nodes.mailbox["eig_d"]
msg = nodes.mailbox["msg"]
degree = msg.size(1)
h = []
for agg in self.aggregators:
if agg.startswith("dir"):
if agg.endswith("av"):
h.append(AGGREGATORS[agg](msg, eig_s, eig_d))
else:
h.append(AGGREGATORS[agg](msg, eig_s, eig_d, h_in))
else:
h.append(AGGREGATORS[agg](msg))
h = torch.cat(h, dim=1)
h = torch.cat(
[
SCALERS[scaler](h, D=degree, delta=self.delta)
if scaler != "identity"
else h
for scaler in self.scalers
],
dim=1,
)
return {"h_neigh": h}
class DGNConv(PNAConv):
r"""Directional Graph Network Layer from `Directional Graph Networks
<https://arxiv.org/abs/2010.02863>`__
DGN introduces two special directional aggregators according to the vector field
:math:`F`, which is defined as the gradient of the low-frequency eigenvectors of graph
laplacian.
The directional average aggregator is defined as
:math:`h_i' = \sum_{j\in\mathcal{N}(i)}\frac{|F_{i,j}|\cdot h_j}{||F_{i,:}||_1+\epsilon}`
The directional derivative aggregator is defined as
:math:`h_i' = \sum_{j\in\mathcal{N}(i)}\frac{F_{i,j}\cdot h_j}{||F_{i,:}||_1+\epsilon}
-h_i\cdot\sum_{j\in\mathcal{N}(i)}\frac{F_{i,j}}{||F_{i,:}||_1+\epsilon}`
:math:`\epsilon` is the infinitesimal to keep the computation numerically stable.
Parameters
----------
in_size : int
Input feature size; i.e. the size of :math:`h_i^l`.
out_size : int
Output feature size; i.e. the size of :math:`h_i^{l+1}`.
aggregators : list of str
List of aggregation function names(each aggregator specifies a way to aggregate
messages from neighbours), selected from:
* ``mean``: the mean of neighbour messages
* ``max``: the maximum of neighbour messages
* ``min``: the minimum of neighbour messages
* ``std``: the standard deviation of neighbour messages
* ``var``: the variance of neighbour messages
* ``sum``: the sum of neighbour messages
* ``moment3``, ``moment4``, ``moment5``: the normalized moments aggregation
:math:`(E[(X-E[X])^n])^{1/n}`
* ``dir{k}-av``: directional average aggregation with directions defined by the k-th
smallest eigenvectors. k can be selected from 1, 2, 3.
* ``dir{k}-dx``: directional derivative aggregation with directions defined by the k-th
smallest eigenvectors. k can be selected from 1, 2, 3.
Note that using directional aggregation requires the LaplacianPE transform on the input
graph for eigenvector computation (the PE size must be >= k above).
scalers: list of str
List of scaler function names, selected from:
* ``identity``: no scaling
* ``amplification``: multiply the aggregated message by :math:`\log(d+1)/\delta`,
where :math:`d` is the in-degree of the node.
* ``attenuation``: multiply the aggregated message by :math:`\delta/\log(d+1)`
delta: float
The in-degree-related normalization factor computed over the training set, used by scalers
for normalization. :math:`E[\log(d+1)]`, where :math:`d` is the in-degree for each node
in the training set.
dropout: float, optional
The dropout ratio. Default: 0.0.
num_towers: int, optional
The number of towers used. Default: 1. Note that in_size and out_size must be divisible
by num_towers.
edge_feat_size: int, optional
The edge feature size. Default: 0.
residual : bool, optional
The bool flag that determines whether to add a residual connection for the
output. Default: True. If in_size and out_size of the DGN conv layer are not
the same, this flag will be set as False forcibly.
Example
-------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import DGNConv
>>> from dgl import LaplacianPE
>>>
>>> # DGN requires precomputed eigenvectors, with 'eig' as feature name.
>>> transform = LaplacianPE(k=3, feat_name='eig')
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> g = transform(g)
>>> eig = g.ndata['eig']
>>> feat = th.ones(6, 10)
>>> conv = DGNConv(10, 10, ['dir1-av', 'dir1-dx', 'sum'], ['identity', 'amplification'], 2.5)
>>> ret = conv(g, feat, eig_vec=eig)
"""
def __init__(
self,
in_size,
out_size,
aggregators,
scalers,
delta,
dropout=0.0,
num_towers=1,
edge_feat_size=0,
residual=True,
):
super(DGNConv, self).__init__(
in_size,
out_size,
aggregators,
scalers,
delta,
dropout,
num_towers,
edge_feat_size,
residual,
)
self.towers = nn.ModuleList(
[
DGNConvTower(
self.tower_in_size,
self.tower_out_size,
aggregators,
scalers,
delta,
dropout=dropout,
edge_feat_size=edge_feat_size,
)
for _ in range(num_towers)
]
)
self.use_eig_vec = False
for aggr in aggregators:
if aggr.startswith("dir"):
self.use_eig_vec = True
break
def forward(self, graph, node_feat, edge_feat=None, eig_vec=None):
r"""
Description
-----------
Compute DGN layer.
Parameters
----------
graph : DGLGraph
The graph.
node_feat : torch.Tensor
The input feature of shape :math:`(N, h_n)`. :math:`N` is the number of
nodes, and :math:`h_n` must be the same as in_size.
edge_feat : torch.Tensor, optional
The edge feature of shape :math:`(M, h_e)`. :math:`M` is the number of
edges, and :math:`h_e` must be the same as edge_feat_size.
eig_vec : torch.Tensor, optional
K smallest non-trivial eigenvectors of Graph Laplacian of shape :math:`(N, K)`.
It is only required when :attr:`aggregators` contains directional aggregators.
Returns
-------
torch.Tensor
The output node feature of shape :math:`(N, h_n')` where :math:`h_n'`
should be the same as out_size.
"""
with graph.local_scope():
if self.use_eig_vec:
graph.ndata["eig"] = eig_vec
return super().forward(graph, node_feat, edge_feat)
python/dgl/nn/pytorch/conv/dotgatconv.py
+242
View File
@@ -0,0 +1,242 @@
"""Torch modules for graph attention networks(GAT)."""
# pylint: disable= no-member, arguments-differ, invalid-name
from torch import nn
from .... import function as fn
from ....base import DGLError
from ....utils import expand_as_pair
from ...functional import edge_softmax
class DotGatConv(nn.Module):
r"""Apply dot product version of self attention in `Graph Attention Network
<https://arxiv.org/pdf/1710.10903.pdf>`__
.. math::
h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} \alpha_{i, j} h_j^{(l)}
where :math:`\alpha_{ij}` is the attention score bewteen node :math:`i` and node :math:`j`:
.. math::
\alpha_{i, j} &= \mathrm{softmax_i}(e_{ij}^{l})
e_{ij}^{l} &= ({W_i^{(l)} h_i^{(l)}})^T \cdot {W_j^{(l)} h_j^{(l)}}
where :math:`W_i` and :math:`W_j` transform node :math:`i`'s and node :math:`j`'s
features into the same dimension, so that when compute note features' similarity,
it can use dot-product.
Parameters
----------
in_feats : int, or pair of ints
Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
DotGatConv 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 head in Multi-Head Attention
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 torch as th
>>> from dgl.nn import DotGatConv
>>> # 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 = th.ones(6, 10)
>>> dotgatconv = DotGatConv(10, 2, num_heads=3)
>>> res = dotgatconv(g, feat)
>>> res
tensor([[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]]], grad_fn=<BinaryReduceBackward>)
>>> # 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 = th.tensor(np.random.rand(2, 5).astype(np.float32))
>>> v_feat = th.tensor(np.random.rand(4, 10).astype(np.float32))
>>> dotgatconv = DotGatConv((5,10), 2, 3)
>>> res = dotgatconv(g, (u_feat, v_feat))
>>> res
tensor([[[-0.6066, 1.0268],
[-0.5945, -0.4801],
[ 0.1594, 0.3825]],
[[ 0.0268, 1.0783],
[ 0.5041, -1.3025],
[ 0.6568, 0.7048]],
[[-0.2688, 1.0543],
[-0.0315, -0.9016],
[ 0.3943, 0.5347]],
[[-0.6066, 1.0268],
[-0.5945, -0.4801],
[ 0.1594, 0.3825]]], grad_fn=<BinaryReduceBackward>)
"""
def __init__(
self, in_feats, out_feats, num_heads, allow_zero_in_degree=False
):
super(DotGatConv, self).__init__()
self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
self._out_feats = out_feats
self._allow_zero_in_degree = allow_zero_in_degree
self._num_heads = num_heads
if isinstance(in_feats, tuple):
self.fc_src = nn.Linear(
self._in_src_feats,
self._out_feats * self._num_heads,
bias=False,
)
self.fc_dst = nn.Linear(
self._in_dst_feats,
self._out_feats * self._num_heads,
bias=False,
)
else:
self.fc = nn.Linear(
self._in_src_feats,
self._out_feats * self._num_heads,
bias=False,
)
def forward(self, graph, feat, get_attention=False):
r"""
Description
-----------
Apply dot product version of self attention in GCN.
Parameters
----------
graph: DGLGraph or bi_partities graph
The graph
feat: torch.Tensor or pair of torch.Tensor
If a torch.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 torch.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
get_attention : bool, optional
Whether to return the attention values. Default to False.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}` is size
of output feature.
torch.Tensor, optional
The attention values of shape :math:`(E, 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``.
"""
graph = graph.local_var()
if not self._allow_zero_in_degree:
if (graph.in_degrees() == 0).any():
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."
)
# check if feat is a tuple
if isinstance(feat, tuple):
h_src = feat[0]
h_dst = feat[1]
feat_src = self.fc_src(h_src).view(
-1, self._num_heads, self._out_feats
)
feat_dst = self.fc_dst(h_dst).view(
-1, self._num_heads, self._out_feats
)
else:
h_src = feat
feat_src = feat_dst = self.fc(h_src).view(
-1, self._num_heads, self._out_feats
)
if graph.is_block:
feat_dst = feat_src[: graph.number_of_dst_nodes()]
# Assign features to nodes
graph.srcdata.update({"ft": feat_src})
graph.dstdata.update({"ft": feat_dst})
# Step 1. dot product
graph.apply_edges(fn.u_dot_v("ft", "ft", "a"))
# Step 2. edge softmax to compute attention scores
graph.edata["sa"] = edge_softmax(
graph, graph.edata["a"] / self._out_feats**0.5
)
# Step 3. Broadcast softmax value to each edge, and aggregate dst node
graph.update_all(
fn.u_mul_e("ft", "sa", "attn"), fn.sum("attn", "agg_u")
)
# output results to the destination nodes
rst = graph.dstdata["agg_u"]
if get_attention:
return rst, graph.edata["sa"]
else:
return rst
python/dgl/nn/pytorch/conv/edgeconv.py
+201
View File
@@ -0,0 +1,201 @@
"""Torch Module for EdgeConv Layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
from torch import nn
from .... import function as fn
from ....base import DGLError
from ....utils import expand_as_pair
class EdgeConv(nn.Module):
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 torch as th
>>> 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 = th.ones(6, 10)
>>> conv = EdgeConv(10, 2)
>>> res = conv(g, feat)
>>> res
tensor([[-0.2347, 0.5849],
[-0.2347, 0.5849],
[-0.2347, 0.5849],
[-0.2347, 0.5849],
[-0.2347, 0.5849],
[-0.2347, 0.5849]], grad_fn=<CopyReduceBackward>)
>>> # 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 = th.rand(2, 5)
>>> v_fea = th.rand(4, 5)
>>> conv = EdgeConv(5, 2, 3)
>>> res = conv(g, (u_fea, v_fea))
>>> res
tensor([[ 1.6375, 0.2085],
[-1.1925, -1.2852],
[ 0.2101, 1.3466],
[ 0.2342, -0.9868]], grad_fn=<CopyReduceBackward>)
"""
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
self.theta = nn.Linear(in_feat, out_feat)
self.phi = nn.Linear(in_feat, out_feat)
if batch_norm:
self.bn = nn.BatchNorm1d(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, feat):
"""
Description
-----------
Forward computation
Parameters
----------
g : DGLGraph
The graph.
feat : Tensor or pair of tensors
:math:`(N, D)` where :math:`N` is the number of nodes and
:math:`D` is the number of feature dimensions.
If a pair of tensors is given, the graph must be a uni-bipartite graph
with only one edge type, and the two tensors must have the same
dimensionality on all except the first axis.
Returns
-------
torch.Tensor
New node features.
Raises
------
DGLError
If there are 0-in-degree nodes in the input graph, it will raise DGLError
since no message will be passed to those nodes. This will cause invalid output.
The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
"""
with g.local_scope():
if not self._allow_zero_in_degree:
if (g.in_degrees() == 0).any():
raise DGLError(
"There are 0-in-degree nodes in the graph, "
"output for those nodes will be invalid. "
"This is harmful for some applications, "
"causing silent performance regression. "
"Adding self-loop on the input graph by "
"calling `g = dgl.add_self_loop(g)` will resolve "
"the issue. Setting ``allow_zero_in_degree`` "
"to be `True` when constructing this module will "
"suppress the check and let the code run."
)
h_src, h_dst = expand_as_pair(feat, g)
g.srcdata["x"] = h_src
g.dstdata["x"] = h_dst
g.apply_edges(fn.v_sub_u("x", "x", "theta"))
g.edata["theta"] = self.theta(g.edata["theta"])
g.dstdata["phi"] = self.phi(g.dstdata["x"])
if not self.batch_norm:
g.update_all(fn.e_add_v("theta", "phi", "e"), fn.max("e", "x"))
else:
g.apply_edges(fn.e_add_v("theta", "phi", "e"))
# Although the official implementation includes a per-edge
# batch norm within EdgeConv, I choose to replace it with a
# global batch norm for a number of reasons:
#
# (1) When the point clouds within each batch do not have the
# same number of points, batch norm would not work.
#
# (2) Even if the point clouds always have the same number of
# points, the points may as well be shuffled even with the
# same (type of) object (and the official implementation
# *does* shuffle the points of the same example for each
# epoch).
#
# For example, the first point of a point cloud of an
# airplane does not always necessarily reside at its nose.
#
# In this case, the learned statistics of each position
# by batch norm is not as meaningful as those learned from
# images.
g.edata["e"] = self.bn(g.edata["e"])
g.update_all(fn.copy_e("e", "e"), fn.max("e", "x"))
return g.dstdata["x"]
python/dgl/nn/pytorch/conv/edgegatconv.py
+390
View File
@@ -0,0 +1,390 @@
"""Torch modules for graph attention networks(GAT)."""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from .... import function as fn
from ....base import DGLError
from ....utils import expand_as_pair
from ...functional import edge_softmax
# pylint: enable=W0235
class EdgeGATConv(nn.Module):
r"""Graph attention layer with edge features from `SCENE
<https://arxiv.org/pdf/2301.03512.pdf>`__
.. math::
\mathbf{v}_i^\prime = \mathbf{\Theta}_\mathrm{s} \cdot \mathbf{v}_i +
\sum\limits_{j \in \mathcal{N}(v_i)} \alpha_{j, i} \left( \mathbf{\Theta}_\mathrm{n}
\cdot \mathbf{v}_j + \mathbf{\Theta}_\mathrm{e} \cdot \mathbf{e}_{j,i} \right)
where :math:`\mathbf{\Theta}` is used to denote learnable weight matrices
for the transformation of features of the node to update (s=self),
neighboring nodes (n=neighbor) and edge features (e=edge).
Attention weights are obtained by
.. math::
\alpha_{j, i} = \mathrm{softmax}_i \Big( \mathrm{LeakyReLU} \big( \mathbf{a}^T
[ \mathbf{\Theta}_\mathrm{n} \cdot \mathbf{v}_i || \mathbf{\Theta}_\mathrm{n}
\cdot \mathbf{v}_j || \mathbf{\Theta}_\mathrm{e} \cdot \mathbf{e}_{j,i} ] \big) \Big)
with :math:`\mathbf{a}` corresponding to a learnable vector.
:math:`\mathrm{softmax_i}` stands for the normalization by all incoming edges of node :math:`i`.
Parameters
----------
in_feats : int, or pair of ints
Input feature size; i.e, the number of dimensions of :math:`\mathbf{v}_i`.
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.
edge_feats: int
Edge feature size; i.e., the number of dimensions of :math:\mathbf{e}_{j,i}`.
out_feats : int
Output feature size; i.e, the number of dimensions of :math:`\mathbf{v}_i^\prime`.
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``.
bias : bool, optional
If True, learns a bias term. Defaults: ``True``.
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 torch as th
>>> from dgl.nn import EdgeGATConv
>>> # Case 1: Homogeneous graph.
>>> num_nodes, num_edges = 8, 30
>>> # Generate a graph.
>>> graph = dgl.rand_graph(num_nodes,num_edges)
>>> node_feats = th.rand((num_nodes, 20))
>>> edge_feats = th.rand((num_edges, 12))
>>> edge_gat = EdgeGATConv(
... in_feats=20,
... edge_feats=12,
... out_feats=15,
... num_heads=3,
... )
>>> # Forward pass.
>>> new_node_feats = edge_gat(graph, node_feats, edge_feats)
>>> new_node_feats.shape
torch.Size([8, 3, 15]) torch.Size([30, 3, 10])
>>> # 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 = th.tensor(np.random.rand(2, 25).astype(np.float32))
>>> v_feat = th.tensor(np.random.rand(4, 30).astype(np.float32))
>>> nfeats = (u_feat,v_feat)
>>> efeats = th.tensor(np.random.rand(5, 15).astype(np.float32))
>>> in_feats = (25,30)
>>> edge_feats = 15
>>> out_feats = 10
>>> num_heads = 3
>>> egat_model = EdgeGATConv(
... in_feats,
... edge_feats,
... out_feats,
... num_heads,
... )
>>> # Forward pass.
>>> new_node_feats, attention_weights = egat_model(g, nfeats, efeats, get_attention=True)
>>> new_node_feats.shape, attention_weights.shape
(torch.Size([4, 3, 10]), torch.Size([5, 3, 1]))
"""
def __init__(
self,
in_feats,
edge_feats,
out_feats,
num_heads,
feat_drop=0.0,
attn_drop=0.0,
negative_slope=0.2,
residual=True,
activation=None,
allow_zero_in_degree=False,
bias=True,
):
super(EdgeGATConv, self).__init__()
self._num_heads = num_heads
self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
self._out_feats = out_feats
self._allow_zero_in_degree = allow_zero_in_degree
if isinstance(in_feats, tuple):
self.fc_src = nn.Linear(
self._in_src_feats, out_feats * num_heads, bias=False
)
self.fc_dst = nn.Linear(
self._in_dst_feats, out_feats * num_heads, bias=False
)
else:
self.fc = nn.Linear(
self._in_src_feats, out_feats * num_heads, bias=False
)
self.attn_l = nn.Parameter(
th.FloatTensor(size=(1, num_heads, out_feats))
)
self.attn_r = nn.Parameter(
th.FloatTensor(size=(1, num_heads, out_feats))
)
self.feat_drop = nn.Dropout(feat_drop)
self.attn_drop = nn.Dropout(attn_drop)
self.leaky_relu = nn.LeakyReLU(negative_slope)
if bias:
self.bias = nn.Parameter(
th.FloatTensor(size=(num_heads * out_feats,))
)
else:
self.register_buffer("bias", None)
if residual:
self.res_fc = nn.Linear(
self._in_dst_feats, num_heads * out_feats, bias=False
)
else:
self.register_buffer("res_fc", None)
self._edge_feats = edge_feats
self.fc_edge = nn.Linear(edge_feats, out_feats * num_heads, bias=False)
self.attn_edge = nn.Parameter(
th.FloatTensor(size=(1, num_heads, out_feats))
)
self.reset_parameters()
self.activation = activation
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The fc weights :math:`\mathbf{\Theta}` are and the
attention weights are using xavier initialization method.
"""
gain = nn.init.calculate_gain("relu")
if hasattr(self, "fc"):
nn.init.xavier_normal_(self.fc.weight, gain=gain)
else:
nn.init.xavier_normal_(self.fc_src.weight, gain=gain)
nn.init.xavier_normal_(self.fc_dst.weight, gain=gain)
nn.init.xavier_normal_(self.attn_l, gain=gain)
nn.init.xavier_normal_(self.attn_r, gain=gain)
nn.init.xavier_normal_(self.fc_edge.weight, gain=gain)
nn.init.xavier_normal_(self.attn_edge, gain=gain)
if self.bias is not None:
nn.init.constant_(self.bias, 0)
if isinstance(self.res_fc, nn.Linear):
nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
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, edge_feat, get_attention=False):
r"""
Description
-----------
Compute graph attention network layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor or pair of torch.Tensor
If a torch.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 torch.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, *, D_{in_{src}})` and :math:`(N_{out}, *, D_{in_{dst}})`.
edge_feat : torch.Tensor
The input edge feature of shape :math:`(E, D_{in_{edge}})`,
where :math:`E` is the number of edges and :math:`D_{in_{edge}}`
the size of the edge features.
get_attention : bool, optional
Whether to return the attention values. Default to False.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, *, H, D_{out})` where :math:`H`
is the number of heads, and :math:`D_{out}` is size of output feature.
torch.Tensor, optional
The attention values of shape :math:`(E, *, H, 1)`. 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() == 0).any():
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]
h_src = self.feat_drop(feat[0])
h_dst = self.feat_drop(feat[1])
if not hasattr(self, "fc_src"):
feat_src = self.fc(h_src).view(
*src_prefix_shape, self._num_heads, self._out_feats
)
feat_dst = self.fc(h_dst).view(
*dst_prefix_shape, self._num_heads, self._out_feats
)
else:
feat_src = self.fc_src(h_src).view(
*src_prefix_shape, self._num_heads, self._out_feats
)
feat_dst = self.fc_dst(h_dst).view(
*dst_prefix_shape, self._num_heads, self._out_feats
)
else:
src_prefix_shape = dst_prefix_shape = feat.shape[:-1]
h_src = h_dst = self.feat_drop(feat)
feat_src = feat_dst = self.fc(h_src).view(
*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:]
# Linearly tranform the edge features.
n_edges = edge_feat.shape[:-1]
feat_edge = self.fc_edge(edge_feat).view(
*n_edges, self._num_heads, self._out_feats
)
# Add edge features to graph.
graph.edata["ft_edge"] = feat_edge
el = (feat_src * self.attn_l).sum(dim=-1).unsqueeze(-1)
er = (feat_dst * self.attn_r).sum(dim=-1).unsqueeze(-1)
# Calculate scalar for each edge.
ee = (feat_edge * self.attn_edge).sum(dim=-1).unsqueeze(-1)
graph.edata["ee"] = ee
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_tmp"))
# e_tmp combines attention weights of source and destination node.
# Add the attention weight of the edge.
graph.edata["e"] = graph.edata["e_tmp"] + graph.edata["ee"]
# Create new edges features that combine the
# features of the source node and the edge features.
graph.apply_edges(fn.u_add_e("ft", "ft_edge", "ft_combined"))
e = self.leaky_relu(graph.edata.pop("e"))
# Compute softmax.
graph.edata["a"] = self.attn_drop(edge_softmax(graph, e))
# For each edge, element-wise multiply the combined features with
# the attention coefficient.
graph.edata["m_combined"] = (
graph.edata["ft_combined"] * graph.edata["a"]
)
# First copy the edge features and then sum them up.
graph.update_all(fn.copy_e("m_combined", "m"), fn.sum("m", "ft"))
rst = graph.dstdata["ft"]
# Residual.
if self.res_fc is not None:
# Use -1 rather than self._num_heads to handle broadcasting.
if h_dst.numel() != 0:
resval = self.res_fc(h_dst).view(
*dst_prefix_shape, -1, self._out_feats
)
rst = rst + resval
# Bias.
if self.bias is not None:
rst = rst + self.bias.view(
*((1,) * len(dst_prefix_shape)),
self._num_heads,
self._out_feats
)
# Activation.
if self.activation:
rst = self.activation(rst)
if get_attention:
return rst, graph.edata["a"]
else:
return rst
python/dgl/nn/pytorch/conv/egatconv.py
+260
View File
@@ -0,0 +1,260 @@
"""Torch modules for graph attention networks with fully valuable edges (EGAT)."""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from torch.nn import init
from .... import function as fn
from ....base import DGLError
from ....utils import expand_as_pair
from ...functional import edge_softmax
# pylint: enable=W0235
class EGATConv(nn.Module):
r"""Graph attention layer that handles edge features from `Rossmann-Toolbox
<https://pubmed.ncbi.nlm.nih.gov/34571541/>`__ (see supplementary data)
The difference lies in how unnormalized attention scores :math:`e_{ij}` are obtained:
.. math::
e_{ij} &= \vec{F} (f_{ij}^{\prime})
f_{ij}^{\prime} &= \mathrm{LeakyReLU}\left(A [ h_{i} \| f_{ij} \| h_{j}]\right)
where :math:`f_{ij}^{\prime}` are edge features, :math:`\mathrm{A}` is weight matrix and
:math:`\vec{F}` is weight vector. After that, resulting node features
:math:`h_{i}^{\prime}` are updated in the same way as in regular GAT.
Parameters
----------
in_node_feats : int, or pair of ints
Input feature size; i.e, the number of dimensions of :math:`h_{i}`.
EGATConv 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.
in_edge_feats : int
Input edge feature size :math:`f_{ij}`.
out_node_feats : int
Output node feature size.
out_edge_feats : int
Output edge feature size :math:`f_{ij}^{\prime}`.
num_heads : int
Number of attention heads.
bias : bool, optional
If True, add bias term to :math:`f_{ij}^{\prime}`. Defaults: ``True``.
Examples
----------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import EGATConv
>>> # Case 1: Homogeneous graph
>>> num_nodes, num_edges = 8, 30
>>> # generate a graph
>>> graph = dgl.rand_graph(num_nodes,num_edges)
>>> node_feats = th.rand((num_nodes, 20))
>>> edge_feats = th.rand((num_edges, 12))
>>> egat = EGATConv(in_node_feats=20,
... in_edge_feats=12,
... out_node_feats=15,
... out_edge_feats=10,
... num_heads=3)
>>> #forward pass
>>> new_node_feats, new_edge_feats = egat(graph, node_feats, edge_feats)
>>> new_node_feats.shape, new_edge_feats.shape
torch.Size([8, 3, 15]) torch.Size([30, 3, 10])
>>> # 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 = th.tensor(np.random.rand(2, 25).astype(np.float32))
>>> v_feat = th.tensor(np.random.rand(4, 30).astype(np.float32))
>>> nfeats = (u_feat,v_feat)
>>> efeats = th.tensor(np.random.rand(5, 15).astype(np.float32))
>>> in_node_feats = (25,30)
>>> in_edge_feats = 15
>>> out_node_feats = 10
>>> out_edge_feats = 5
>>> num_heads = 3
>>> egat_model = EGATConv(in_node_feats,
... in_edge_feats,
... out_node_feats,
... out_edge_feats,
... num_heads,
... bias=True)
>>> #forward pass
>>> new_node_feats,
>>> new_edge_feats,
>>> attentions = egat_model(g, nfeats, efeats, get_attention=True)
>>> new_node_feats.shape, new_edge_feats.shape, attentions.shape
(torch.Size([4, 3, 10]), torch.Size([5, 3, 5]), torch.Size([5, 3, 1]))
"""
def __init__(
self,
in_node_feats,
in_edge_feats,
out_node_feats,
out_edge_feats,
num_heads,
bias=True,
):
super().__init__()
self._num_heads = num_heads
self._in_src_node_feats, self._in_dst_node_feats = expand_as_pair(
in_node_feats
)
self._out_node_feats = out_node_feats
self._out_edge_feats = out_edge_feats
if isinstance(in_node_feats, tuple):
self.fc_node_src = nn.Linear(
self._in_src_node_feats, out_node_feats * num_heads, bias=False
)
self.fc_ni = nn.Linear(
self._in_src_node_feats, out_edge_feats * num_heads, bias=False
)
self.fc_nj = nn.Linear(
self._in_dst_node_feats, out_edge_feats * num_heads, bias=False
)
else:
self.fc_node_src = nn.Linear(
self._in_src_node_feats, out_node_feats * num_heads, bias=False
)
self.fc_ni = nn.Linear(
self._in_src_node_feats, out_edge_feats * num_heads, bias=False
)
self.fc_nj = nn.Linear(
self._in_src_node_feats, out_edge_feats * num_heads, bias=False
)
self.fc_fij = nn.Linear(
in_edge_feats, out_edge_feats * num_heads, bias=False
)
self.attn = nn.Parameter(
th.FloatTensor(size=(1, num_heads, out_edge_feats))
)
if bias:
self.bias = nn.Parameter(
th.FloatTensor(size=(num_heads * out_edge_feats,))
)
else:
self.register_buffer("bias", None)
self.reset_parameters()
def reset_parameters(self):
"""
Reinitialize learnable parameters.
"""
gain = init.calculate_gain("relu")
init.xavier_normal_(self.fc_node_src.weight, gain=gain)
init.xavier_normal_(self.fc_ni.weight, gain=gain)
init.xavier_normal_(self.fc_fij.weight, gain=gain)
init.xavier_normal_(self.fc_nj.weight, gain=gain)
init.xavier_normal_(self.attn, gain=gain)
init.constant_(self.bias, 0)
def forward(
self, graph, nfeats, efeats, edge_weight=None, get_attention=False
):
r"""
Compute new node and edge features.
Parameters
----------
graph : DGLGraph
The graph.
nfeat : torch.Tensor or pair of torch.Tensor
If a torch.Tensor is given, the input feature of shape :math:`(N, D_{in})`
where:
:math:`D_{in}` is size of input node feature,
:math:`N` is the number of nodes.
If a pair of torch.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, D_{in_{src}})` and
:math:`(N_{out}, D_{in_{dst}})`.
efeats: torch.Tensor
The input edge feature of shape :math:`(E, F_{in})`
where:
:math:`F_{in}` is size of input node feature,
:math:`E` is the number of edges.
edge_weight : torch.Tensor, optional
A 1D tensor of edge weight values. Shape: :math:`(|E|,)`.
get_attention : bool, optional
Whether to return the attention values. Default to False.
Returns
-------
pair of torch.Tensor
node output features followed by edge output features.
The node output feature is of shape :math:`(N, H, D_{out})`
The edge output feature is of shape :math:`(F, H, F_{out})`
where:
:math:`H` is the number of heads,
:math:`D_{out}` is size of output node feature,
:math:`F_{out}` is size of output edge feature.
torch.Tensor, optional
The attention values of shape :math:`(E, H, 1)`.
This is returned only when :attr:`get_attention` is ``True``.
"""
with graph.local_scope():
if (graph.in_degrees() == 0).any():
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."
)
# calc edge attention
# same trick way as in dgl.nn.pytorch.GATConv, but also includes edge feats
# https://github.com/dmlc/dgl/blob/master/python/dgl/nn/pytorch/conv/gatconv.py
if isinstance(nfeats, tuple):
nfeats_src, nfeats_dst = nfeats
else:
nfeats_src = nfeats_dst = nfeats
f_ni = self.fc_ni(nfeats_src)
f_nj = self.fc_nj(nfeats_dst)
f_fij = self.fc_fij(efeats)
graph.srcdata.update({"f_ni": f_ni})
graph.dstdata.update({"f_nj": f_nj})
# add ni, nj factors
graph.apply_edges(fn.u_add_v("f_ni", "f_nj", "f_tmp"))
# add fij to node factor
f_out = graph.edata.pop("f_tmp") + f_fij
if self.bias is not None:
f_out = f_out + self.bias
f_out = nn.functional.leaky_relu(f_out)
f_out = f_out.view(-1, self._num_heads, self._out_edge_feats)
# compute attention factor
e = (f_out * self.attn).sum(dim=-1).unsqueeze(-1)
graph.edata["a"] = edge_softmax(graph, e)
if edge_weight is not None:
graph.edata["a"] = graph.edata["a"] * edge_weight.tile(
1, self._num_heads, 1
).transpose(0, 2)
graph.srcdata["h_out"] = self.fc_node_src(nfeats_src).view(
-1, self._num_heads, self._out_node_feats
)
# calc weighted sum
graph.update_all(
fn.u_mul_e("h_out", "a", "m"), fn.sum("m", "h_out")
)
h_out = graph.dstdata["h_out"].view(
-1, self._num_heads, self._out_node_feats
)
if get_attention:
return h_out, f_out, graph.edata.pop("a")
else:
return h_out, f_out
python/dgl/nn/pytorch/conv/egnnconv.py
+163
View File
@@ -0,0 +1,163 @@
"""Torch Module for E(n) Equivariant Graph Convolutional Layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch
import torch.nn as nn
from .... import function as fn
class EGNNConv(nn.Module):
r"""Equivariant Graph Convolutional Layer from `E(n) Equivariant Graph
Neural Networks <https://arxiv.org/abs/2102.09844>`__
.. math::
m_{ij}=\phi_e(h_i^l, h_j^l, ||x_i^l-x_j^l||^2, a_{ij})
x_i^{l+1} = x_i^l + C\sum_{j\in\mathcal{N}(i)}(x_i^l-x_j^l)\phi_x(m_{ij})
m_i = \sum_{j\in\mathcal{N}(i)} m_{ij}
h_i^{l+1} = \phi_h(h_i^l, m_i)
where :math:`h_i`, :math:`x_i`, :math:`a_{ij}` are node features, coordinate
features, and edge features respectively. :math:`\phi_e`, :math:`\phi_h`, and
:math:`\phi_x` are two-layer MLPs. :math:`C` is a constant for normalization,
computed as :math:`1/|\mathcal{N}(i)|`.
Parameters
----------
in_size : int
Input feature size; i.e. the size of :math:`h_i^l`.
hidden_size : int
Hidden feature size; i.e. the size of hidden layer in the two-layer MLPs in
:math:`\phi_e, \phi_x, \phi_h`.
out_size : int
Output feature size; i.e. the size of :math:`h_i^{l+1}`.
edge_feat_size : int, optional
Edge feature size; i.e. the size of :math:`a_{ij}`. Default: 0.
Example
-------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import EGNNConv
>>>
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> node_feat, coord_feat, edge_feat = th.ones(6, 10), th.ones(6, 3), th.ones(6, 2)
>>> conv = EGNNConv(10, 10, 10, 2)
>>> h, x = conv(g, node_feat, coord_feat, edge_feat)
"""
def __init__(self, in_size, hidden_size, out_size, edge_feat_size=0):
super(EGNNConv, self).__init__()
self.in_size = in_size
self.hidden_size = hidden_size
self.out_size = out_size
self.edge_feat_size = edge_feat_size
act_fn = nn.SiLU()
# \phi_e
self.edge_mlp = nn.Sequential(
# +1 for the radial feature: ||x_i - x_j||^2
nn.Linear(in_size * 2 + edge_feat_size + 1, hidden_size),
act_fn,
nn.Linear(hidden_size, hidden_size),
act_fn,
)
# \phi_h
self.node_mlp = nn.Sequential(
nn.Linear(in_size + hidden_size, hidden_size),
act_fn,
nn.Linear(hidden_size, out_size),
)
# \phi_x
self.coord_mlp = nn.Sequential(
nn.Linear(hidden_size, hidden_size),
act_fn,
nn.Linear(hidden_size, 1, bias=False),
)
def message(self, edges):
"""message function for EGNN"""
# concat features for edge mlp
if self.edge_feat_size > 0:
f = torch.cat(
[
edges.src["h"],
edges.dst["h"],
edges.data["radial"],
edges.data["a"],
],
dim=-1,
)
else:
f = torch.cat(
[edges.src["h"], edges.dst["h"], edges.data["radial"]], dim=-1
)
msg_h = self.edge_mlp(f)
msg_x = self.coord_mlp(msg_h) * edges.data["x_diff"]
return {"msg_x": msg_x, "msg_h": msg_h}
def forward(self, graph, node_feat, coord_feat, edge_feat=None):
r"""
Description
-----------
Compute EGNN layer.
Parameters
----------
graph : DGLGraph
The graph.
node_feat : torch.Tensor
The input feature of shape :math:`(N, h_n)`. :math:`N` is the number of
nodes, and :math:`h_n` must be the same as in_size.
coord_feat : torch.Tensor
The coordinate feature of shape :math:`(N, h_x)`. :math:`N` is the
number of nodes, and :math:`h_x` can be any positive integer.
edge_feat : torch.Tensor, optional
The edge feature of shape :math:`(M, h_e)`. :math:`M` is the number of
edges, and :math:`h_e` must be the same as edge_feat_size.
Returns
-------
node_feat_out : torch.Tensor
The output node feature of shape :math:`(N, h_n')` where :math:`h_n'`
is the same as out_size.
coord_feat_out: torch.Tensor
The output coordinate feature of shape :math:`(N, h_x)` where :math:`h_x`
is the same as the input coordinate feature dimension.
"""
with graph.local_scope():
# node feature
graph.ndata["h"] = node_feat
# coordinate feature
graph.ndata["x"] = coord_feat
# edge feature
if self.edge_feat_size > 0:
assert edge_feat is not None, "Edge features must be provided."
graph.edata["a"] = edge_feat
# get coordinate diff & radial features
graph.apply_edges(fn.u_sub_v("x", "x", "x_diff"))
graph.edata["radial"] = (
graph.edata["x_diff"].square().sum(dim=1).unsqueeze(-1)
)
# normalize coordinate difference
graph.edata["x_diff"] = graph.edata["x_diff"] / (
graph.edata["radial"].sqrt() + 1e-30
)
graph.apply_edges(self.message)
graph.update_all(fn.copy_e("msg_x", "m"), fn.mean("m", "x_neigh"))
graph.update_all(fn.copy_e("msg_h", "m"), fn.sum("m", "h_neigh"))
h_neigh, x_neigh = graph.ndata["h_neigh"], graph.ndata["x_neigh"]
h = self.node_mlp(torch.cat([node_feat, h_neigh], dim=-1))
x = coord_feat + x_neigh
return h, x
python/dgl/nn/pytorch/conv/gatconv.py
+370
View File
@@ -0,0 +1,370 @@
"""Torch modules for graph attention networks(GAT)."""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch 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 Identity
# pylint: enable=W0235
class GATConv(nn.Module):
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``.
bias : bool, optional
If True, learns a bias term. Defaults: ``True``.
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 torch as th
>>> 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 = th.ones(6, 10)
>>> gatconv = GATConv(10, 2, num_heads=3)
>>> res = gatconv(g, feat)
>>> res
tensor([[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]],
[[ 3.4570, 1.8634],
[ 1.3805, -0.0762],
[ 1.0390, -1.1479]]], grad_fn=<BinaryReduceBackward>)
>>> # 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 = th.tensor(np.random.rand(2, 5).astype(np.float32))
>>> v_feat = th.tensor(np.random.rand(4, 10).astype(np.float32))
>>> gatconv = GATConv((5,10), 2, 3)
>>> res = gatconv(g, (u_feat, v_feat))
>>> res
tensor([[[-0.6066, 1.0268],
[-0.5945, -0.4801],
[ 0.1594, 0.3825]],
[[ 0.0268, 1.0783],
[ 0.5041, -1.3025],
[ 0.6568, 0.7048]],
[[-0.2688, 1.0543],
[-0.0315, -0.9016],
[ 0.3943, 0.5347]],
[[-0.6066, 1.0268],
[-0.5945, -0.4801],
[ 0.1594, 0.3825]]], grad_fn=<BinaryReduceBackward>)
"""
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,
bias=True,
):
super(GATConv, self).__init__()
self._num_heads = num_heads
self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
self._out_feats = out_feats
self._allow_zero_in_degree = allow_zero_in_degree
if isinstance(in_feats, tuple):
self.fc_src = nn.Linear(
self._in_src_feats, out_feats * num_heads, bias=False
)
self.fc_dst = nn.Linear(
self._in_dst_feats, out_feats * num_heads, bias=False
)
else:
self.fc = nn.Linear(
self._in_src_feats, out_feats * num_heads, bias=False
)
self.attn_l = nn.Parameter(
th.FloatTensor(size=(1, num_heads, out_feats))
)
self.attn_r = nn.Parameter(
th.FloatTensor(size=(1, num_heads, out_feats))
)
self.feat_drop = nn.Dropout(feat_drop)
self.attn_drop = nn.Dropout(attn_drop)
self.leaky_relu = nn.LeakyReLU(negative_slope)
self.has_linear_res = False
self.has_explicit_bias = False
if residual:
if self._in_dst_feats != out_feats * num_heads:
self.res_fc = nn.Linear(
self._in_dst_feats, num_heads * out_feats, bias=bias
)
self.has_linear_res = True
else:
self.res_fc = Identity()
else:
self.register_buffer("res_fc", None)
if bias and not self.has_linear_res:
self.bias = nn.Parameter(
th.FloatTensor(size=(num_heads * out_feats,))
)
self.has_explicit_bias = True
else:
self.register_buffer("bias", None)
self.reset_parameters()
self.activation = activation
def reset_parameters(self):
"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The fc weights :math:`W^{(l)}` are initialized using Glorot uniform initialization.
The attention weights are using xavier initialization method.
"""
gain = nn.init.calculate_gain("relu")
if hasattr(self, "fc"):
nn.init.xavier_normal_(self.fc.weight, gain=gain)
else:
nn.init.xavier_normal_(self.fc_src.weight, gain=gain)
nn.init.xavier_normal_(self.fc_dst.weight, gain=gain)
nn.init.xavier_normal_(self.attn_l, gain=gain)
nn.init.xavier_normal_(self.attn_r, gain=gain)
if self.has_explicit_bias:
nn.init.constant_(self.bias, 0)
if isinstance(self.res_fc, nn.Linear):
nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
if self.res_fc.bias is not None:
nn.init.constant_(self.res_fc.bias, 0)
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, edge_weight=None, get_attention=False):
r"""
Description
-----------
Compute graph attention network layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor or pair of torch.Tensor
If a torch.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 torch.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, *, D_{in_{src}})` and :math:`(N_{out}, *, D_{in_{dst}})`.
edge_weight : torch.Tensor, optional
A 1D tensor of edge weight values. Shape: :math:`(|E|,)`.
get_attention : bool, optional
Whether to return the attention values. Default to False.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, *, H, D_{out})` where :math:`H`
is the number of heads, and :math:`D_{out}` is size of output feature.
torch.Tensor, optional
The attention values of shape :math:`(E, *, H, 1)`, where :math:`E` is the number of
edges. This is returned only when :attr:`get_attention` is ``True``.
Raises
------
DGLError
If there are 0-in-degree nodes in the input graph, it will raise DGLError
since no message will be passed to those nodes. This will cause invalid output.
The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
"""
with graph.local_scope():
if not self._allow_zero_in_degree:
if (graph.in_degrees() == 0).any():
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]
h_src = self.feat_drop(feat[0])
h_dst = self.feat_drop(feat[1])
if not hasattr(self, "fc_src"):
feat_src = self.fc(h_src).view(
*src_prefix_shape, self._num_heads, self._out_feats
)
feat_dst = self.fc(h_dst).view(
*dst_prefix_shape, self._num_heads, self._out_feats
)
else:
feat_src = self.fc_src(h_src).view(
*src_prefix_shape, self._num_heads, self._out_feats
)
feat_dst = self.fc_dst(h_dst).view(
*dst_prefix_shape, self._num_heads, self._out_feats
)
else:
src_prefix_shape = dst_prefix_shape = feat.shape[:-1]
h_src = h_dst = self.feat_drop(feat)
feat_src = feat_dst = self.fc(h_src).view(
*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).sum(dim=-1).unsqueeze(-1)
er = (feat_dst * self.attn_r).sum(dim=-1).unsqueeze(-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))
if edge_weight is not None:
graph.edata["a"] = graph.edata["a"] * edge_weight.tile(
1, self._num_heads, 1
).transpose(0, 2)
# message passing
graph.update_all(fn.u_mul_e("ft", "a", "m"), fn.sum("m", "ft"))
rst = graph.dstdata["ft"]
# residual
if self.res_fc is not None:
# Use -1 rather than self._num_heads to handle broadcasting
if h_dst.numel() != 0:
resval = self.res_fc(h_dst).view(
*dst_prefix_shape, -1, self._out_feats
)
rst = rst + resval
# bias
if self.has_explicit_bias:
rst = rst + self.bias.view(
*((1,) * len(dst_prefix_shape)),
self._num_heads,
self._out_feats
)
# activation
if self.activation:
rst = self.activation(rst)
if get_attention:
return rst, graph.edata["a"]
else:
return rst
python/dgl/nn/pytorch/conv/gatedgcnconv.py
+173
View File
@@ -0,0 +1,173 @@
"""Torch Module for GatedGCN layer"""
# pylint: disable= no-member, arguments-differ, invalid-name, cell-var-from-loop
import torch
import torch.nn.functional as F
from torch import nn
from .... import function as fn
class GatedGCNConv(nn.Module):
r"""Gated graph convolutional layer from `Benchmarking Graph Neural Networks
<https://arxiv.org/abs/2003.00982>`__
.. math::
e_{ij}^{l+1}=D^l h_{i}^{l}+E^l h_{j}^{l}+C^l e_{ij}^{l}
norm_{ij}=\Sigma_{j\in N_{i}} \sigma\left(e_{ij}^{l+1}\right)+\varepsilon
\hat{e}_{ij}^{l+1}=\sigma(e_{ij}^{l+1}) / norm_{ij}
h_{i}^{l+1}=A^l h_{i}^{l}+\Sigma_{j \in N_{i}} \hat{e}_{ij}^{l+1} \odot B^l h_{j}^{l}
where :math:`h_{i}^{l}` is node :math:`i` feature of layer :math:`l`,
:math:`e_{ij}^{l}` is edge :math:`ij` feature of layer :math:`l`,
:math:`\sigma` is sigmoid function, :math:`\varepsilon` is a small fixed constant
for numerical stability, :math:`A^l, B^l, C^l, D^l, E^l` are linear layers.
Parameters
----------
input_feats : int
Input feature size; i.e, the number of dimensions of :math:`h_{i}^{l}`.
edge_feats: int
Edge feature size; i.e., the number of dimensions of :math:`e_{ij}^{l}`.
output_feats : int
Output feature size; i.e., the number of dimensions of :math:`h_{i}^{l+1}`.
dropout : float, optional
Dropout rate on node and edge feature. Default: ``0``.
batch_norm : bool, optional
Whether to include batch normalization on node and edge feature. Default: ``True``.
residual : bool, optional
Whether to include residual connections. Default: ``True``.
activation : callable activation function/layer or None, optional
If not None, apply an activation function to the updated node features.
Default: ``F.relu``.
Example
-------
>>> import dgl
>>> import torch as th
>>> import torch.nn.functional as F
>>> from dgl.nn import GatedGCNConv
>>> num_nodes, num_edges = 8, 30
>>> graph = dgl.rand_graph(num_nodes,num_edges)
>>> node_feats = th.rand(num_nodes, 20)
>>> edge_feats = th.rand(num_edges, 12)
>>> gatedGCN = GatedGCNConv(20, 12, 20)
>>> new_node_feats, new_edge_feats = gatedGCN(graph, node_feats, edge_feats)
>>> new_node_feats.shape, new_edge_feats.shape
(torch.Size([8, 20]), torch.Size([30, 20]))
"""
def __init__(
self,
input_feats,
edge_feats,
output_feats,
dropout=0,
batch_norm=True,
residual=True,
activation=F.relu,
):
super(GatedGCNConv, self).__init__()
self.dropout = nn.Dropout(dropout)
self.batch_norm = batch_norm
self.residual = residual
if input_feats != output_feats or edge_feats != output_feats:
self.residual = False
# Linearly transform the node features.
self.A = nn.Linear(input_feats, output_feats, bias=True)
self.B = nn.Linear(input_feats, output_feats, bias=True)
self.D = nn.Linear(input_feats, output_feats, bias=True)
self.E = nn.Linear(input_feats, output_feats, bias=True)
# Linearly transform the edge features.
self.C = nn.Linear(edge_feats, output_feats, bias=True)
# Batch normalization on the node/edge features.
self.bn_node = nn.BatchNorm1d(output_feats)
self.bn_edge = nn.BatchNorm1d(output_feats)
self.activation = activation
def forward(self, graph, feat, edge_feat):
"""
Description
-----------
Compute gated graph convolution layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor
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.
edge_feat : torch.Tensor
The input edge feature of shape :math:`(E, D_{edge})`,
where :math:`E` is the number of edges and :math:`D_{edge}`
is the size of the edge features.
Returns
-------
torch.Tensor
The output node feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
is the output feature size.
torch.Tensor
The output edge feature of shape :math:`(E, D_{out})` where :math:`D_{out}`
is the output feature size.
"""
with graph.local_scope():
# For residual connection
h_in = feat
e_in = edge_feat
graph.ndata["Ah"] = self.A(feat)
graph.ndata["Bh"] = self.B(feat)
graph.ndata["Dh"] = self.D(feat)
graph.ndata["Eh"] = self.E(feat)
graph.edata["Ce"] = self.C(edge_feat)
graph.apply_edges(fn.u_add_v("Dh", "Eh", "DEh"))
# Get edge feature
graph.edata["e"] = graph.edata["DEh"] + graph.edata["Ce"]
graph.edata["sigma"] = torch.sigmoid(graph.edata["e"])
graph.update_all(
fn.u_mul_e("Bh", "sigma", "m"), fn.sum("m", "sum_sigma_h")
)
graph.update_all(fn.copy_e("sigma", "m"), fn.sum("m", "sum_sigma"))
graph.ndata["h"] = graph.ndata["Ah"] + graph.ndata[
"sum_sigma_h"
] / (graph.ndata["sum_sigma"] + 1e-6)
# Result of graph convolution.
feat = graph.ndata["h"]
edge_feat = graph.edata["e"]
# Batch normalization.
if self.batch_norm:
feat = self.bn_node(feat)
edge_feat = self.bn_edge(edge_feat)
# Non-linear activation.
if self.activation:
feat = self.activation(feat)
edge_feat = self.activation(edge_feat)
# Residual connection.
if self.residual:
feat = h_in + feat
edge_feat = e_in + edge_feat
feat = self.dropout(feat)
edge_feat = self.dropout(edge_feat)
return feat, edge_feat
python/dgl/nn/pytorch/conv/gatedgraphconv.py
+173
View File
@@ -0,0 +1,173 @@
"""Torch Module for Gated Graph Convolution layer"""
# pylint: disable= no-member, arguments-differ, invalid-name, cell-var-from-loop
import torch as th
from torch import nn
from torch.nn import init
from .... import function as fn
class GatedGraphConv(nn.Module):
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``.
Example
-------
>>> import dgl
>>> import numpy as np
>>> import torch as th
>>> from dgl.nn import GatedGraphConv
>>>
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 10)
>>> conv = GatedGraphConv(10, 10, 2, 3)
>>> etype = th.tensor([0,1,2,0,1,2])
>>> res = conv(g, feat, etype)
>>> res
tensor([[ 0.4652, 0.4458, 0.5169, 0.4126, 0.4847, 0.2303, 0.2757, 0.7721,
0.0523, 0.0857],
[ 0.0832, 0.1388, -0.5643, 0.7053, -0.2524, -0.3847, 0.7587, 0.8245,
0.9315, 0.4063],
[ 0.6340, 0.4096, 0.7692, 0.2125, 0.2106, 0.4542, -0.0580, 0.3364,
-0.1376, 0.4948],
[ 0.5551, 0.7946, 0.6220, 0.8058, 0.5711, 0.3063, -0.5454, 0.2272,
-0.6931, -0.1607],
[ 0.2644, 0.2469, -0.6143, 0.6008, -0.1516, -0.3781, 0.5878, 0.7993,
0.9241, 0.1835],
[ 0.6393, 0.3447, 0.3893, 0.4279, 0.3342, 0.3809, 0.0406, 0.5030,
0.1342, 0.0425]], grad_fn=<AddBackward0>)
"""
def __init__(self, in_feats, out_feats, n_steps, n_etypes, bias=True):
super(GatedGraphConv, self).__init__()
assert in_feats <= out_feats, "out_feats must be not less than in_feats"
self._in_feats = in_feats
self._out_feats = out_feats
self._n_steps = n_steps
self._n_etypes = n_etypes
self.linears = nn.ModuleList(
[nn.Linear(out_feats, out_feats) for _ in range(n_etypes)]
)
self.gru = nn.GRUCell(out_feats, out_feats, bias=bias)
self.reset_parameters()
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The model parameters are initialized using Glorot uniform initialization
and the bias is initialized to be zero.
"""
gain = init.calculate_gain("relu")
self.gru.reset_parameters()
for linear in self.linears:
init.xavier_normal_(linear.weight, gain=gain)
init.zeros_(linear.bias)
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, etypes=None):
"""
Description
-----------
Compute Gated Graph Convolution layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor
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, or None
The edge type tensor of shape :math:`(E,)` where :math:`E` is
the number of edges of the graph. When there's only one edge type,
this argument can be skipped
Returns
-------
torch.Tensor
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"
)
if self._n_etypes != 1:
assert (
etypes.min() >= 0 and etypes.max() < self._n_etypes
), "edge type indices out of range [0, {})".format(
self._n_etypes
)
zero_pad = feat.new_zeros(
(feat.shape[0], self._out_feats - feat.shape[1])
)
feat = th.cat([feat, zero_pad], -1)
for _ in range(self._n_steps):
if self._n_etypes == 1 and etypes is None:
# Fast path when graph has only one edge type
graph.ndata["h"] = self.linears[0](feat)
graph.update_all(fn.copy_u("h", "m"), fn.sum("m", "a"))
a = graph.ndata.pop("a") # (N, D)
else:
graph.ndata["h"] = feat
for i in range(self._n_etypes):
eids = (
th.nonzero(etypes == i, as_tuple=False)
.view(-1)
.type(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") # (N, D)
feat = self.gru(a, feat)
return feat
python/dgl/nn/pytorch/conv/gatv2conv.py
+335
View File
@@ -0,0 +1,335 @@
"""Torch modules for graph attention networks v2 (GATv2)."""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch 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 Identity
# pylint: enable=W0235
class GATv2Conv(nn.Module):
r"""GATv2 from `How Attentive are Graph Attention Networks?
<https://arxiv.org/pdf/2105.14491.pdf>`__
.. math::
h_i^{(l+1)} = \sum_{j\in \mathcal{N}(i)} \alpha_{ij}^{(l)} W^{(l)}_{right} 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)} &= {\vec{a}^T}^{(l)}\mathrm{LeakyReLU}\left(
W^{(l)}_{left} h_{i} + W^{(l)}_{right} 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)}`.
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``.
bias : bool, optional
If set to :obj:`False`, the layer will not learn
an additive bias. (default: :obj:`True`)
share_weights : bool, optional
If set to :obj:`True`, the same matrix for :math:`W_{left}` and :math:`W_{right}` in
the above equations, will be applied to the source and the target node of every edge.
(default: :obj:`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 applied 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 torch as th
>>> from dgl.nn import GATv2Conv
>>> # 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 = th.ones(6, 10)
>>> gatv2conv = GATv2Conv(10, 2, num_heads=3)
>>> res = gatv2conv(g, feat)
>>> res
tensor([[[ 1.9599, 1.0239],
[ 3.2015, -0.5512],
[ 2.3700, -2.2182]],
[[ 1.9599, 1.0239],
[ 3.2015, -0.5512],
[ 2.3700, -2.2182]],
[[ 1.9599, 1.0239],
[ 3.2015, -0.5512],
[ 2.3700, -2.2182]],
[[ 1.9599, 1.0239],
[ 3.2015, -0.5512],
[ 2.3700, -2.2182]],
[[ 1.9599, 1.0239],
[ 3.2015, -0.5512],
[ 2.3700, -2.2182]],
[[ 1.9599, 1.0239],
[ 3.2015, -0.5512],
[ 2.3700, -2.2182]]], grad_fn=<GSpMMBackward>)
>>> # 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 = th.tensor(np.random.rand(2, 5).astype(np.float32))
>>> v_feat = th.tensor(np.random.rand(4, 10).astype(np.float32))
>>> gatv2conv = GATv2Conv((5,10), 2, 3)
>>> res = gatv2conv(g, (u_feat, v_feat))
>>> res
tensor([[[-0.0935, -0.4273],
[-1.1850, 0.1123],
[-0.2002, 0.1155]],
[[ 0.1908, -1.2095],
[-0.0129, 0.6408],
[-0.8135, 0.1157]],
[[ 0.0596, -0.8487],
[-0.5421, 0.4022],
[-0.4805, 0.1156]],
[[-0.0935, -0.4273],
[-1.1850, 0.1123],
[-0.2002, 0.1155]]], grad_fn=<GSpMMBackward>)
"""
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,
bias=True,
share_weights=False,
):
super(GATv2Conv, self).__init__()
self._num_heads = num_heads
self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
self._out_feats = out_feats
self._allow_zero_in_degree = allow_zero_in_degree
if isinstance(in_feats, tuple):
self.fc_src = nn.Linear(
self._in_src_feats, out_feats * num_heads, bias=bias
)
self.fc_dst = nn.Linear(
self._in_dst_feats, out_feats * num_heads, bias=bias
)
else:
self.fc_src = nn.Linear(
self._in_src_feats, out_feats * num_heads, bias=bias
)
if share_weights:
self.fc_dst = self.fc_src
else:
self.fc_dst = nn.Linear(
self._in_src_feats, out_feats * num_heads, bias=bias
)
self.attn = nn.Parameter(th.FloatTensor(size=(1, num_heads, out_feats)))
self.feat_drop = nn.Dropout(feat_drop)
self.attn_drop = nn.Dropout(attn_drop)
self.leaky_relu = nn.LeakyReLU(negative_slope)
if residual:
if self._in_dst_feats != out_feats * num_heads:
self.res_fc = nn.Linear(
self._in_dst_feats, num_heads * out_feats, bias=bias
)
else:
self.res_fc = Identity()
else:
self.register_buffer("res_fc", None)
self.activation = activation
self.share_weights = share_weights
self.bias = bias
self.reset_parameters()
def reset_parameters(self):
"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The fc weights :math:`W^{(l)}` are initialized using Glorot uniform initialization.
The attention weights are using xavier initialization method.
"""
gain = nn.init.calculate_gain("relu")
nn.init.xavier_normal_(self.fc_src.weight, gain=gain)
if self.bias:
nn.init.constant_(self.fc_src.bias, 0)
if not self.share_weights:
nn.init.xavier_normal_(self.fc_dst.weight, gain=gain)
if self.bias:
nn.init.constant_(self.fc_dst.bias, 0)
nn.init.xavier_normal_(self.attn, gain=gain)
if isinstance(self.res_fc, nn.Linear):
nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
if self.bias:
nn.init.constant_(self.res_fc.bias, 0)
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 : torch.Tensor or pair of torch.Tensor
If a torch.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 torch.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
get_attention : bool, optional
Whether to return the attention values. Default to False.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, H, D_{out})` where :math:`H`
is the number of heads, and :math:`D_{out}` is size of output feature.
torch.Tensor, optional
The attention values of shape :math:`(E, H, 1)`, where :math:`E` is the number of
edges. This is returned only when :attr:`get_attention` is ``True``.
Raises
------
DGLError
If there are 0-in-degree nodes in the input graph, it will raise DGLError
since no message will be passed to those nodes. This will cause invalid output.
The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
"""
with graph.local_scope():
if not self._allow_zero_in_degree:
if (graph.in_degrees() == 0).any():
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):
h_src = self.feat_drop(feat[0])
h_dst = self.feat_drop(feat[1])
feat_src = self.fc_src(h_src).view(
-1, self._num_heads, self._out_feats
)
feat_dst = self.fc_dst(h_dst).view(
-1, self._num_heads, self._out_feats
)
else:
h_src = h_dst = self.feat_drop(feat)
feat_src = self.fc_src(h_src).view(
-1, self._num_heads, self._out_feats
)
if self.share_weights:
feat_dst = feat_src
else:
feat_dst = self.fc_dst(h_dst).view(
-1, self._num_heads, self._out_feats
)
if graph.is_block:
feat_dst = feat_dst[: graph.number_of_dst_nodes()]
h_dst = h_dst[: graph.number_of_dst_nodes()]
graph.srcdata.update(
{"el": feat_src}
) # (num_src_edge, num_heads, out_dim)
graph.dstdata.update({"er": feat_dst})
graph.apply_edges(fn.u_add_v("el", "er", "e"))
e = self.leaky_relu(
graph.edata.pop("e")
) # (num_src_edge, num_heads, out_dim)
e = (
(e * self.attn).sum(dim=-1).unsqueeze(dim=2)
) # (num_edge, num_heads, 1)
# compute softmax
graph.edata["a"] = self.attn_drop(
edge_softmax(graph, e)
) # (num_edge, num_heads)
# message passing
graph.update_all(fn.u_mul_e("el", "a", "m"), fn.sum("m", "ft"))
rst = graph.dstdata["ft"]
# residual
if self.res_fc is not None:
if h_dst.numel() != 0:
resval = self.res_fc(h_dst).view(
h_dst.shape[0], -1, self._out_feats
)
rst = rst + resval
# activation
if self.activation:
rst = self.activation(rst)
if get_attention:
return rst, graph.edata["a"]
else:
return rst
python/dgl/nn/pytorch/conv/gcn2conv.py
+285
View File
@@ -0,0 +1,285 @@
"""Torch Module for Graph Convolutional Network via Initial residual
and Identity mapping (GCNII) layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import math
import torch as th
from torch import nn
from .... import function as fn
from ....base import DGLError
from .graphconv import EdgeWeightNorm
class GCN2Conv(nn.Module):
r"""Graph Convolutional Network via Initial residual
and Identity mapping (GCNII) from `Simple and Deep Graph Convolutional
Networks <https://arxiv.org/abs/2007.02133>`__
It is mathematically is defined as follows:
.. math::
\mathbf{h}^{(l+1)} =\left( (1 - \alpha)(\mathbf{D}^{-1/2} \mathbf{\hat{A}}
\mathbf{D}^{-1/2})\mathbf{h}^{(l)} + \alpha {\mathbf{h}^{(0)}} \right)
\left( (1 - \beta_l) \mathbf{I} + \beta_l \mathbf{W} \right)
where :math:`\mathbf{\hat{A}}` is the adjacency matrix with self-loops,
:math:`\mathbf{D}_{ii} = \sum_{j=0} \mathbf{A}_{ij}` is its diagonal degree matrix,
:math:`\mathbf{h}^{(0)}` is the initial node features,
:math:`\mathbf{h}^{(l)}` is the feature of layer :math:`l`,
:math:`\alpha` is the fraction of initial node features, and
:math:`\beta_l` is the hyperparameter to tune the strength of identity mapping.
It is defined by :math:`\beta_l = \log(\frac{\lambda}{l}+1)\approx\frac{\lambda}{l}`,
where :math:`\lambda` is a hyperparameter. :math:`\beta` ensures that the decay of
the weight matrix adaptively increases as we stack more layers.
Parameters
----------
in_feats : int
Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
layer : int
the index of current layer.
alpha : float
The fraction of the initial input features. Default: ``0.1``
lambda_ : float
The hyperparameter to ensure the decay of the weight matrix
adaptively increases. Default: ``1``
project_initial_features : bool
Whether to share a weight matrix between initial features and
smoothed features. Default: ``True``
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``.
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 torch as th
>>> from dgl.nn import GCN2Conv
>>> # Homogeneous graph
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 3)
>>> g = dgl.add_self_loop(g)
>>> conv1 = GCN2Conv(3, layer=1, alpha=0.5, \
... project_initial_features=True, allow_zero_in_degree=True)
>>> conv2 = GCN2Conv(3, layer=2, alpha=0.5, \
... project_initial_features=True, allow_zero_in_degree=True)
>>> res = feat
>>> res = conv1(g, res, feat)
>>> res = conv2(g, res, feat)
>>> print(res)
tensor([[1.3803, 3.3191, 2.9572],
[1.3803, 3.3191, 2.9572],
[1.3803, 3.3191, 2.9572],
[1.4770, 3.8326, 3.2451],
[1.3623, 3.2102, 2.8679],
[1.3803, 3.3191, 2.9572]], grad_fn=<AddBackward0>)
"""
def __init__(
self,
in_feats,
layer,
alpha=0.1,
lambda_=1,
project_initial_features=True,
allow_zero_in_degree=False,
bias=True,
activation=None,
):
super().__init__()
self._in_feats = in_feats
self._project_initial_features = project_initial_features
self.alpha = alpha
self.beta = math.log(lambda_ / layer + 1)
self._bias = bias
self._activation = activation
self._allow_zero_in_degree = allow_zero_in_degree
self.weight1 = nn.Parameter(th.Tensor(self._in_feats, self._in_feats))
if self._project_initial_features:
self.register_parameter("weight2", None)
else:
self.weight2 = nn.Parameter(
th.Tensor(self._in_feats, self._in_feats)
)
if self._bias:
self.bias = nn.Parameter(th.Tensor(self._in_feats))
else:
self.register_parameter("bias", None)
self.reset_parameters()
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
"""
nn.init.normal_(self.weight1)
if not self._project_initial_features:
nn.init.normal_(self.weight2)
if self._bias:
nn.init.zeros_(self.bias)
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, feat_0, edge_weight=None):
r"""
Description
-----------
Compute graph convolution.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor
The input feature of shape
:math:`(N, D_{in})`
where :math:`D_{in}` is the size of input feature and :math:`N` is the number of nodes.
feat_0 : torch.Tensor
The initial feature of shape :math:`(N, D_{in})`
edge_weight: torch.Tensor, optional
edge_weight to use in the message passing process. This is equivalent to
using weighted adjacency matrix in the equation above, and
:math:`\tilde{D}^{-1/2}\tilde{A} \tilde{D}^{-1/2}`
is based on :class:`dgl.nn.pytorch.conv.graphconv.EdgeWeightNorm`.
Returns
-------
torch.Tensor
The output feature
Raises
------
DGLError
If there are 0-in-degree nodes in the input graph, it will raise DGLError
since no message will be passed to those nodes. This will cause invalid output.
The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
Note
----
* Input shape: :math:`(N, *, \text{in_feats})` where * means any number of additional
dimensions, :math:`N` is the number of nodes.
* Output shape: :math:`(N, *, \text{out_feats})` where all but the last dimension are
the same shape as the input.
* Weight shape: :math:`(\text{in_feats}, \text{out_feats})`.
"""
with graph.local_scope():
if not self._allow_zero_in_degree:
if (graph.in_degrees() == 0).any():
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."
)
# normalize to get smoothed representation
if edge_weight is None:
degs = graph.in_degrees().to(feat).clamp(min=1)
norm = th.pow(degs, -0.5)
norm = norm.to(feat.device).unsqueeze(1)
else:
edge_weight = EdgeWeightNorm("both")(graph, edge_weight)
if edge_weight is None:
feat = feat * norm
graph.ndata["h"] = feat
msg_func = fn.copy_u("h", "m")
if edge_weight is not None:
graph.edata["_edge_weight"] = edge_weight
msg_func = fn.u_mul_e("h", "_edge_weight", "m")
graph.update_all(msg_func, fn.sum("m", "h"))
feat = graph.ndata.pop("h")
if edge_weight is None:
feat = feat * norm
# scale
feat = feat * (1 - self.alpha)
# initial residual connection to the first layer
feat_0 = feat_0[: feat.size(0)] * self.alpha
feat_sum = feat + feat_0
if self._project_initial_features:
feat_proj_sum = feat_sum @ self.weight1
else:
feat_proj_sum = feat @ self.weight1 + feat_0 @ self.weight2
rst = (1 - self.beta) * feat_sum + self.beta * feat_proj_sum
if self._bias:
rst = rst + self.bias
if self._activation is not None:
rst = self._activation(rst)
return rst
def extra_repr(self):
"""Set the extra representation of the module,
which will come into effect when printing the model.
"""
summary = "in={_in_feats}"
summary += ", alpha={alpha}, beta={beta}"
if "self._bias" in self.__dict__:
summary += ", bias={bias}"
if "self._activation" in self.__dict__:
summary += ", activation={_activation}"
return summary.format(**self.__dict__)
python/dgl/nn/pytorch/conv/ginconv.py
+158
View File
@@ -0,0 +1,158 @@
"""Torch Module for Graph Isomorphism Network layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from .... import function as fn
from ....utils import expand_as_pair
class GINConv(nn.Module):
r"""Graph Isomorphism Network layer from `How Powerful are Graph
Neural Networks? <https://arxiv.org/pdf/1810.00826.pdf>`__
.. math::
h_i^{(l+1)} = f_\Theta \left((1 + \epsilon) h_i^{l} +
\mathrm{aggregate}\left(\left\{h_j^{l}, j\in\mathcal{N}(i)
\right\}\right)\right)
If a weight tensor on each edge is provided, the weighted graph convolution is defined as:
.. math::
h_i^{(l+1)} = f_\Theta \left((1 + \epsilon) h_i^{l} +
\mathrm{aggregate}\left(\left\{e_{ji} h_j^{l}, j\in\mathcal{N}(i)
\right\}\right)\right)
where :math:`e_{ji}` is the weight on the edge from node :math:`j` to node :math:`i`.
Please make sure that `e_{ji}` is broadcastable with `h_j^{l}`.
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, default: None.
aggregator_type : str
Aggregator type to use (``sum``, ``max`` or ``mean``), default: 'sum'.
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``.
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 torch as th
>>> from dgl.nn import GINConv
>>>
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 10)
>>> lin = th.nn.Linear(10, 10)
>>> conv = GINConv(lin, 'max')
>>> res = conv(g, feat)
>>> res
tensor([[-0.4821, 0.0207, -0.7665, 0.5721, -0.4682, -0.2134, -0.5236, 1.2855,
0.8843, -0.8764],
[-0.4821, 0.0207, -0.7665, 0.5721, -0.4682, -0.2134, -0.5236, 1.2855,
0.8843, -0.8764],
[-0.4821, 0.0207, -0.7665, 0.5721, -0.4682, -0.2134, -0.5236, 1.2855,
0.8843, -0.8764],
[-0.4821, 0.0207, -0.7665, 0.5721, -0.4682, -0.2134, -0.5236, 1.2855,
0.8843, -0.8764],
[-0.4821, 0.0207, -0.7665, 0.5721, -0.4682, -0.2134, -0.5236, 1.2855,
0.8843, -0.8764],
[-0.1804, 0.0758, -0.5159, 0.3569, -0.1408, -0.1395, -0.2387, 0.7773,
0.5266, -0.4465]], grad_fn=<AddmmBackward>)
>>> # With activation
>>> from torch.nn.functional import relu
>>> conv = GINConv(lin, 'max', activation=relu)
>>> res = conv(g, feat)
>>> res
tensor([[5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
0.0000],
[5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
0.0000],
[5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
0.0000],
[5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
0.0000],
[5.0118, 0.0000, 0.0000, 3.9091, 1.3371, 0.0000, 0.0000, 0.0000, 0.0000,
0.0000],
[2.5011, 0.0000, 0.0089, 2.0541, 0.8262, 0.0000, 0.0000, 0.1371, 0.0000,
0.0000]], grad_fn=<ReluBackward0>)
"""
def __init__(
self,
apply_func=None,
aggregator_type="sum",
init_eps=0,
learn_eps=False,
activation=None,
):
super(GINConv, self).__init__()
self.apply_func = apply_func
self._aggregator_type = aggregator_type
self.activation = activation
if aggregator_type not in ("sum", "max", "mean"):
raise KeyError(
"Aggregator type {} not recognized.".format(aggregator_type)
)
# to specify whether eps is trainable or not.
if learn_eps:
self.eps = th.nn.Parameter(th.FloatTensor([init_eps]))
else:
self.register_buffer("eps", th.FloatTensor([init_eps]))
def forward(self, graph, feat, edge_weight=None):
r"""
Description
-----------
Compute Graph Isomorphism Network layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor or pair of torch.Tensor
If a torch.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 torch.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, D_{in})` and :math:`(N_{out}, D_{in})`.
If ``apply_func`` is not None, :math:`D_{in}` should
fit the input dimensionality requirement of ``apply_func``.
edge_weight : torch.Tensor, optional
Optional tensor on the edge. If given, the convolution will weight
with regard to the message.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, D_{out})` where
:math:`D_{out}` is the output dimensionality of ``apply_func``.
If ``apply_func`` is None, :math:`D_{out}` should be the same
as input dimensionality.
"""
_reducer = getattr(fn, self._aggregator_type)
with graph.local_scope():
aggregate_fn = fn.copy_u("h", "m")
if edge_weight is not None:
assert edge_weight.shape[0] == graph.num_edges()
graph.edata["_edge_weight"] = edge_weight
aggregate_fn = fn.u_mul_e("h", "_edge_weight", "m")
feat_src, feat_dst = expand_as_pair(feat, graph)
graph.srcdata["h"] = feat_src
graph.update_all(aggregate_fn, _reducer("m", "neigh"))
rst = (1 + self.eps) * feat_dst + graph.dstdata["neigh"]
if self.apply_func is not None:
rst = self.apply_func(rst)
# activation
if self.activation is not None:
rst = self.activation(rst)
return rst
python/dgl/nn/pytorch/conv/gineconv.py
+97
View File
@@ -0,0 +1,97 @@
"""Torch Module for Graph Isomorphism Network layer variant with edge features"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
import torch.nn.functional as F
from torch import nn
from .... import function as fn
from ....utils import expand_as_pair
class GINEConv(nn.Module):
r"""Graph Isomorphism Network with Edge Features, introduced by
`Strategies for Pre-training Graph Neural Networks <https://arxiv.org/abs/1905.12265>`__
.. math::
h_i^{(l+1)} = f_\Theta \left((1 + \epsilon) h_i^{l} +
\sum_{j\in\mathcal{N}(i)}\mathrm{ReLU}(h_j^{l} + e_{j,i}^{l})\right)
where :math:`e_{j,i}^{l}` is the edge feature.
Parameters
----------
apply_func : callable module or None
The :math:`f_\Theta` in the formula. If not None, it will be applied to
the updated node features. The default value is None.
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``.
Examples
--------
>>> import dgl
>>> import torch
>>> import torch.nn as nn
>>> from dgl.nn import GINEConv
>>> g = dgl.graph(([0, 1, 2], [1, 1, 3]))
>>> in_feats = 10
>>> out_feats = 20
>>> nfeat = torch.randn(g.num_nodes(), in_feats)
>>> efeat = torch.randn(g.num_edges(), in_feats)
>>> conv = GINEConv(nn.Linear(in_feats, out_feats))
>>> res = conv(g, nfeat, efeat)
>>> print(res.shape)
torch.Size([4, 20])
"""
def __init__(self, apply_func=None, init_eps=0, learn_eps=False):
super(GINEConv, self).__init__()
self.apply_func = apply_func
# to specify whether eps is trainable or not.
if learn_eps:
self.eps = nn.Parameter(th.FloatTensor([init_eps]))
else:
self.register_buffer("eps", th.FloatTensor([init_eps]))
def message(self, edges):
r"""User-defined Message Function"""
return {"m": F.relu(edges.src["hn"] + edges.data["he"])}
def forward(self, graph, node_feat, edge_feat):
r"""Forward computation.
Parameters
----------
graph : DGLGraph
The graph.
node_feat : torch.Tensor or pair of torch.Tensor
If a torch.Tensor is given, it is the input feature of shape :math:`(N, D_{in})` where
:math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
If a pair of torch.Tensor is given, 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 feature size requirement of ``apply_func``.
edge_feat : torch.Tensor
Edge feature. It is a tensor of shape :math:`(E, D_{in})` where :math:`E`
is the number of edges.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, D_{out})` where
:math:`D_{out}` is the output feature size of ``apply_func``.
If ``apply_func`` is None, :math:`D_{out}` should be the same
as :math:`D_{in}`.
"""
with graph.local_scope():
feat_src, feat_dst = expand_as_pair(node_feat, graph)
graph.srcdata["hn"] = feat_src
graph.edata["he"] = edge_feat
graph.update_all(self.message, fn.sum("m", "neigh"))
rst = (1 + self.eps) * feat_dst + graph.dstdata["neigh"]
if self.apply_func is not None:
rst = self.apply_func(rst)
return rst
python/dgl/nn/pytorch/conv/gmmconv.py
+268
View File
@@ -0,0 +1,268 @@
"""Torch Module for GMM Conv"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from torch.nn import init
from .... import function as fn
from ....base import DGLError
from ....utils import expand_as_pair
from ..utils import Identity
class GMMConv(nn.Module):
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 torch as th
>>> 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 = th.ones(6, 10)
>>> conv = GMMConv(10, 2, 3, 2, 'mean')
>>> pseudo = th.ones(12, 3)
>>> res = conv(g, feat, pseudo)
>>> res
tensor([[-0.3462, -0.2654],
[-0.3462, -0.2654],
[-0.3462, -0.2654],
[-0.3462, -0.2654],
[-0.3462, -0.2654],
[-0.3462, -0.2654]], grad_fn=<AddBackward0>)
>>> # 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 = th.rand(2, 5)
>>> v_fea = th.rand(4, 10)
>>> pseudo = th.ones(5, 3)
>>> conv = GMMConv((10, 5), 2, 3, 2, 'mean')
>>> res = conv(g, (u_fea, v_fea), pseudo)
>>> res
tensor([[-0.1107, -0.1559],
[-0.1646, -0.2326],
[-0.1377, -0.1943],
[-0.1107, -0.1559]], grad_fn=<AddBackward0>)
"""
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)
)
self.mu = nn.Parameter(th.Tensor(n_kernels, dim))
self.inv_sigma = nn.Parameter(th.Tensor(n_kernels, dim))
self.fc = nn.Linear(
self._in_src_feats, n_kernels * out_feats, bias=False
)
if residual:
if self._in_dst_feats != out_feats:
self.res_fc = nn.Linear(
self._in_dst_feats, out_feats, bias=False
)
else:
self.res_fc = Identity()
else:
self.register_buffer("res_fc", None)
if bias:
self.bias = nn.Parameter(th.Tensor(out_feats))
else:
self.register_buffer("bias", None)
self.reset_parameters()
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The fc parameters are initialized using Glorot uniform initialization
and the bias is initialized to be zero.
The mu weight is initialized using normal distribution and
inv_sigma is initialized with constant value 1.0.
"""
gain = init.calculate_gain("relu")
init.xavier_normal_(self.fc.weight, gain=gain)
if isinstance(self.res_fc, nn.Linear):
init.xavier_normal_(self.res_fc.weight, gain=gain)
init.normal_(self.mu.data, 0, 0.1)
init.constant_(self.inv_sigma.data, 1)
if self.bias is not None:
init.zeros_(self.bias.data)
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 : torch.Tensor
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 : torch.Tensor
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
-------
torch.Tensor
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``.
"""
with graph.local_scope():
if not self._allow_zero_in_degree:
if (graph.in_degrees() == 0).any():
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"] = self.fc(feat_src).view(
-1, self._n_kernels, self._out_feats
)
E = graph.num_edges()
# compute gaussian weight
gaussian = -0.5 * (
(
pseudo.view(E, 1, self._dim)
- self.mu.view(1, self._n_kernels, self._dim)
)
** 2
)
gaussian = gaussian * (
self.inv_sigma.view(1, self._n_kernels, self._dim) ** 2
)
gaussian = th.exp(gaussian.sum(dim=-1, keepdim=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
return rst
python/dgl/nn/pytorch/conv/graphconv.py
+488
View File
@@ -0,0 +1,488 @@
"""Torch modules for graph convolutions(GCN)."""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from torch.nn import init
from .... import function as fn
from ....base import DGLError
from ....convert import block_to_graph
from ....heterograph import DGLBlock
from ....transforms import reverse
from ....utils import expand_as_pair
class EdgeWeightNorm(nn.Module):
r"""This module normalizes positive scalar edge weights on a graph
following the form in `GCN <https://arxiv.org/abs/1609.02907>`__.
Mathematically, setting ``norm='both'`` yields the following normalization term:
.. math::
c_{ji} = (\sqrt{\sum_{k\in\mathcal{N}(j)}e_{jk}}\sqrt{\sum_{k\in\mathcal{N}(i)}e_{ki}})
And, setting ``norm='right'`` yields the following normalization term:
.. math::
c_{ji} = (\sum_{k\in\mathcal{N}(i)}e_{ki})
where :math:`e_{ji}` is the scalar weight on the edge from node :math:`j` to node :math:`i`.
The module returns the normalized weight :math:`e_{ji} / c_{ji}`.
Parameters
----------
norm : str, optional
The normalizer as specified above. Default is `'both'`.
eps : float, optional
A small offset value in the denominator. Default is 0.
Examples
--------
>>> import dgl
>>> import numpy as np
>>> import torch as th
>>> from dgl.nn import EdgeWeightNorm, GraphConv
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> g = dgl.add_self_loop(g)
>>> feat = th.ones(6, 10)
>>> edge_weight = th.tensor([0.5, 0.6, 0.4, 0.7, 0.9, 0.1, 1, 1, 1, 1, 1, 1])
>>> norm = EdgeWeightNorm(norm='both')
>>> norm_edge_weight = norm(g, edge_weight)
>>> conv = GraphConv(10, 2, norm='none', weight=True, bias=True)
>>> res = conv(g, feat, edge_weight=norm_edge_weight)
>>> print(res)
tensor([[-1.1849, -0.7525],
[-1.3514, -0.8582],
[-1.2384, -0.7865],
[-1.9949, -1.2669],
[-1.3658, -0.8674],
[-0.8323, -0.5286]], grad_fn=<AddBackward0>)
"""
def __init__(self, norm="both", eps=0.0):
super(EdgeWeightNorm, self).__init__()
self._norm = norm
self._eps = eps
def forward(self, graph, edge_weight):
r"""
Description
-----------
Compute normalized edge weight for the GCN model.
Parameters
----------
graph : DGLGraph
The graph.
edge_weight : torch.Tensor
Unnormalized scalar weights on the edges.
The shape is expected to be :math:`(|E|)`.
Returns
-------
torch.Tensor
The normalized edge weight.
Raises
------
DGLError
Case 1:
The edge weight is multi-dimensional. Currently this module
only supports a scalar weight on each edge.
Case 2:
The edge weight has non-positive values with ``norm='both'``.
This will trigger square root and division by a non-positive number.
"""
with graph.local_scope():
if isinstance(graph, DGLBlock):
graph = block_to_graph(graph)
if len(edge_weight.shape) > 1:
raise DGLError(
"Currently the normalization is only defined "
"on scalar edge weight. Please customize the "
"normalization for your high-dimensional weights."
)
if self._norm == "both" and th.any(edge_weight <= 0).item():
raise DGLError(
'Non-positive edge weight detected with `norm="both"`. '
"This leads to square root of zero or negative values."
)
dev = graph.device
dtype = edge_weight.dtype
graph.srcdata["_src_out_w"] = th.ones(
graph.number_of_src_nodes(), dtype=dtype, device=dev
)
graph.dstdata["_dst_in_w"] = th.ones(
graph.number_of_dst_nodes(), dtype=dtype, device=dev
)
graph.edata["_edge_w"] = edge_weight
if self._norm == "both":
reversed_g = reverse(graph)
reversed_g.edata["_edge_w"] = edge_weight
reversed_g.update_all(
fn.copy_e("_edge_w", "m"), fn.sum("m", "out_weight")
)
degs = reversed_g.dstdata["out_weight"] + self._eps
norm = th.pow(degs, -0.5)
graph.srcdata["_src_out_w"] = norm
if self._norm != "none":
graph.update_all(
fn.copy_e("_edge_w", "m"), fn.sum("m", "in_weight")
)
degs = graph.dstdata["in_weight"] + self._eps
if self._norm == "both":
norm = th.pow(degs, -0.5)
else:
norm = 1.0 / degs
graph.dstdata["_dst_in_w"] = norm
graph.apply_edges(
lambda e: {
"_norm_edge_weights": e.src["_src_out_w"]
* e.dst["_dst_in_w"]
* e.data["_edge_w"]
}
)
return graph.edata["_norm_edge_weights"]
# pylint: disable=W0235
class GraphConv(nn.Module):
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_{ji}}h_j^{(l)}W^{(l)})
where :math:`\mathcal{N}(i)` is the set of neighbors of node :math:`i`,
:math:`c_{ji}` is the product of the square root of node degrees
(i.e., :math:`c_{ji} = \sqrt{|\mathcal{N}(j)|}\sqrt{|\mathcal{N}(i)|}`),
and :math:`\sigma` is an activation function.
If a weight tensor on each edge is provided, the weighted graph convolution is defined as:
.. math::
h_i^{(l+1)} = \sigma(b^{(l)} + \sum_{j\in\mathcal{N}(i)}\frac{e_{ji}}{c_{ji}}h_j^{(l)}W^{(l)})
where :math:`e_{ji}` is the scalar weight on the edge from node :math:`j` to node :math:`i`.
This is NOT equivalent to the weighted graph convolutional network formulation in the paper.
To customize the normalization term :math:`c_{ji}`, one can first set ``norm='none'`` for
the model, and send the pre-normalized :math:`e_{ji}` to the forward computation. We provide
:class:`~dgl.nn.pytorch.EdgeWeightNorm` to normalize scalar edge weight following the GCN paper.
Parameters
----------
in_feats : int
Input feature size; i.e, the number of dimensions of :math:`h_j^{(l)}`.
out_feats : int
Output feature size; i.e., the number of dimensions of :math:`h_i^{(l+1)}`.
norm : str, optional
How to apply the normalizer. Can be one of the following values:
* ``right``, to divide the aggregated messages by each node's in-degrees,
which is equivalent to averaging the received messages.
* ``none``, where no normalization is applied.
* ``both`` (default), where the messages are scaled with :math:`1/c_{ji}` above, equivalent
to symmetric normalization.
* ``left``, to divide the messages sent out from each node by its out-degrees,
equivalent to random walk normalization.
weight : bool, optional
If True, apply a linear layer. Otherwise, aggregating the messages
without a weight matrix.
bias : bool, optional
If True, adds a learnable bias to the output. Default: ``True``.
activation : callable activation function/layer or None, optional
If not None, applies an activation function to the updated node features.
Default: ``None``.
allow_zero_in_degree : bool, optional
If there are 0-in-degree nodes in the graph, output for those nodes will be invalid
since no message will be passed to those nodes. This is harmful for some applications
causing silent performance regression. This module will raise a DGLError if it detects
0-in-degree nodes in input graph. By setting ``True``, it will suppress the check
and let the users handle it by themselves. Default: ``False``.
Attributes
----------
weight : torch.Tensor
The learnable weight tensor.
bias : torch.Tensor
The learnable bias tensor.
Note
----
Zero in-degree nodes will lead to invalid output value. This is because no message
will be passed to those nodes, the aggregation function will be appied on empty input.
A common practice to avoid this is to add a self-loop for each node in the graph if
it is homogeneous, which can be achieved by:
>>> g = ... # a DGLGraph
>>> g = dgl.add_self_loop(g)
Calling ``add_self_loop`` will not work for some graphs, for example, heterogeneous graph
since the edge type can not be decided for self_loop edges. Set ``allow_zero_in_degree``
to ``True`` for those cases to unblock the code and handle zero-in-degree nodes manually.
A common practise to handle this is to filter out the nodes with zero-in-degree when use
after conv.
Examples
--------
>>> import dgl
>>> import numpy as np
>>> import torch as th
>>> 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 = th.ones(6, 10)
>>> conv = GraphConv(10, 2, norm='both', weight=True, bias=True)
>>> res = conv(g, feat)
>>> print(res)
tensor([[ 1.3326, -0.2797],
[ 1.4673, -0.3080],
[ 1.3326, -0.2797],
[ 1.6871, -0.3541],
[ 1.7711, -0.3717],
[ 1.0375, -0.2178]], grad_fn=<AddBackward0>)
>>> # 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)
tensor([[-0.2473, -0.4631],
[-0.3497, -0.6549],
[-0.3497, -0.6549],
[-0.4221, -0.7905],
[-0.3497, -0.6549],
[ 0.0000, 0.0000]], grad_fn=<AddBackward0>)
>>> # Case 2: Unidirectional bipartite graph
>>> u = [0, 1, 0, 0, 1]
>>> v = [0, 1, 2, 3, 2]
>>> g = dgl.heterograph({('_U', '_E', '_V') : (u, v)})
>>> u_fea = th.rand(2, 5)
>>> v_fea = th.rand(4, 5)
>>> conv = GraphConv(5, 2, norm='both', weight=True, bias=True)
>>> res = conv(g, (u_fea, v_fea))
>>> res
tensor([[-0.2994, 0.6106],
[-0.4482, 0.5540],
[-0.5287, 0.8235],
[-0.2994, 0.6106]], grad_fn=<AddBackward0>)
"""
def __init__(
self,
in_feats,
out_feats,
norm="both",
weight=True,
bias=True,
activation=None,
allow_zero_in_degree=False,
):
super(GraphConv, self).__init__()
if norm not in ("none", "both", "right", "left"):
raise DGLError(
'Invalid norm value. Must be either "none", "both", "right" or "left".'
' But got "{}".'.format(norm)
)
self._in_feats = in_feats
self._out_feats = out_feats
self._norm = norm
self._allow_zero_in_degree = allow_zero_in_degree
if weight:
self.weight = nn.Parameter(th.Tensor(in_feats, out_feats))
else:
self.register_parameter("weight", None)
if bias:
self.bias = nn.Parameter(th.Tensor(out_feats))
else:
self.register_parameter("bias", None)
self.reset_parameters()
self._activation = activation
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The model parameters are initialized as in the
`original implementation <https://github.com/tkipf/gcn/blob/master/gcn/layers.py>`__
where the weight :math:`W^{(l)}` is initialized using Glorot uniform initialization
and the bias is initialized to be zero.
"""
if self.weight is not None:
init.xavier_uniform_(self.weight)
if self.bias is not None:
init.zeros_(self.bias)
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, edge_weight=None):
r"""
Description
-----------
Compute graph convolution.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor or pair of torch.Tensor
If a torch.Tensor is given, it represents the input feature of shape
:math:`(N, D_{in})`
where :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
If a pair of torch.Tensor is given, which is the case for bipartite graph, the pair
must contain two tensors of shape :math:`(N_{in}, D_{in_{src}})` and
:math:`(N_{out}, D_{in_{dst}})`.
weight : torch.Tensor, optional
Optional external weight tensor.
edge_weight : torch.Tensor, optional
Optional tensor on the edge. If given, the convolution will weight
with regard to the message.
Returns
-------
torch.Tensor
The output feature
Raises
------
DGLError
Case 1:
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``.
Case 2:
External weight is provided while at the same time the module
has defined its own weight parameter.
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() == 0).any():
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."
)
aggregate_fn = fn.copy_u("h", "m")
if edge_weight is not None:
assert edge_weight.shape[0] == graph.num_edges()
graph.edata["_edge_weight"] = edge_weight
aggregate_fn = fn.u_mul_e("h", "_edge_weight", "m")
# (BarclayII) For RGCN on heterogeneous graphs we need to support GCN on bipartite.
feat_src, feat_dst = expand_as_pair(feat, graph)
if self._norm in ["left", "both"]:
degs = graph.out_degrees().to(feat_src).clamp(min=1)
if self._norm == "both":
norm = th.pow(degs, -0.5)
else:
norm = 1.0 / degs
shp = norm.shape + (1,) * (feat_src.dim() - 1)
norm = th.reshape(norm, shp)
feat_src = feat_src * norm
if weight is not None:
if self.weight is not None:
raise DGLError(
"External weight is provided while at the same time the"
" module has defined its own weight parameter. Please"
" create the module with flag weight=False."
)
else:
weight = self.weight
if self._in_feats > self._out_feats:
# mult W first to reduce the feature size for aggregation.
if weight is not None:
feat_src = th.matmul(feat_src, weight)
graph.srcdata["h"] = feat_src
graph.update_all(aggregate_fn, fn.sum(msg="m", out="h"))
rst = graph.dstdata["h"]
else:
# aggregate first then mult W
graph.srcdata["h"] = feat_src
graph.update_all(aggregate_fn, fn.sum(msg="m", out="h"))
rst = graph.dstdata["h"]
if weight is not None:
rst = th.matmul(rst, weight)
if self._norm in ["right", "both"]:
degs = graph.in_degrees().to(feat_dst).clamp(min=1)
if self._norm == "both":
norm = th.pow(degs, -0.5)
else:
norm = 1.0 / degs
shp = norm.shape + (1,) * (feat_dst.dim() - 1)
norm = th.reshape(norm, shp)
rst = rst * norm
if self.bias is not None:
rst = rst + self.bias
if self._activation is not None:
rst = self._activation(rst)
return rst
def extra_repr(self):
"""Set the extra representation of the module,
which will come into effect when printing the model.
"""
summary = "in={_in_feats}, out={_out_feats}"
summary += ", normalization={_norm}"
if "_activation" in self.__dict__:
summary += ", activation={_activation}"
return summary.format(**self.__dict__)
python/dgl/nn/pytorch/conv/grouprevres.py
+257
View File
@@ -0,0 +1,257 @@
"""Torch module for grouped reversible residual connections for GNNs"""
# pylint: disable= no-member, arguments-differ, invalid-name, C0116, R1728
from copy import deepcopy
import numpy as np
import torch
import torch.nn as nn
class InvertibleCheckpoint(torch.autograd.Function):
r"""Extension of torch.autograd"""
@staticmethod
def forward(ctx, fn, fn_inverse, num_inputs, *inputs_and_weights):
ctx.fn = fn
ctx.fn_inverse = fn_inverse
ctx.weights = inputs_and_weights[num_inputs:]
inputs = inputs_and_weights[:num_inputs]
ctx.input_requires_grad = []
with torch.no_grad():
# Make a detached copy, which shares the storage
x = []
for element in inputs:
if isinstance(element, torch.Tensor):
x.append(element.detach())
ctx.input_requires_grad.append(element.requires_grad)
else:
x.append(element)
ctx.input_requires_grad.append(None)
# Detach the output, which then allows discarding the intermediary results
outputs = ctx.fn(*x).detach_()
# clear memory of input node features
inputs[1].untyped_storage().resize_(0)
# store for backward pass
ctx.inputs = [inputs]
ctx.outputs = [outputs]
return outputs
@staticmethod
def backward(ctx, *grad_outputs):
if not torch.autograd._is_checkpoint_valid():
raise RuntimeError(
"InvertibleCheckpoint is not compatible with .grad(), \
please use .backward() if possible"
)
# retrieve input and output tensor nodes
if len(ctx.outputs) == 0:
raise RuntimeError(
"Trying to perform backward on the InvertibleCheckpoint \
for more than once."
)
inputs = ctx.inputs.pop()
outputs = ctx.outputs.pop()
# reconstruct input node features
with torch.no_grad():
# inputs[0] is DGLGraph and inputs[1] is input node features
inputs_inverted = ctx.fn_inverse(
*((inputs[0], outputs) + inputs[2:])
)
# clear memory of outputs
outputs.untyped_storage().resize_(0)
x = inputs[1]
x.untyped_storage().resize_(int(np.prod(x.size())))
x.set_(inputs_inverted)
# compute gradients
with torch.set_grad_enabled(True):
detached_inputs = []
for i, element in enumerate(inputs):
if isinstance(element, torch.Tensor):
element = element.detach()
element.requires_grad = ctx.input_requires_grad[i]
detached_inputs.append(element)
detached_inputs = tuple(detached_inputs)
temp_output = ctx.fn(*detached_inputs)
filtered_detached_inputs = tuple(
filter(
lambda x: getattr(x, "requires_grad", False), detached_inputs
)
)
gradients = torch.autograd.grad(
outputs=(temp_output,),
inputs=filtered_detached_inputs + ctx.weights,
grad_outputs=grad_outputs,
)
input_gradients = []
i = 0
for rg in ctx.input_requires_grad:
if rg:
input_gradients.append(gradients[i])
i += 1
else:
input_gradients.append(None)
gradients = tuple(input_gradients) + gradients[-len(ctx.weights) :]
return (None, None, None) + gradients
class GroupRevRes(nn.Module):
r"""Grouped reversible residual connections for GNNs, as introduced in
`Training Graph Neural Networks with 1000 Layers <https://arxiv.org/abs/2106.07476>`__
It uniformly partitions an input node feature :math:`X` into :math:`C` groups
:math:`X_1, X_2, \cdots, X_C` across the channel dimension. Besides, it makes
:math:`C` copies of the input GNN module :math:`f_{w1}, \cdots, f_{wC}`. In the
forward pass, each GNN module only takes the corresponding group of node features.
The output node representations :math:`X^{'}` are computed as follows.
.. math::
X_0^{'} = \sum_{i=2}^{C}X_i
X_i^{'} = f_{wi}(X_{i-1}^{'}, g, U) + X_i, i\in\{1,\cdots,C\}
X^{'} = X_1^{'} \, \Vert \, \ldots \, \Vert \, X_C^{'}
where :math:`g` is the input graph, :math:`U` is arbitrary additional input arguments like
edge features, and :math:`\, \Vert \,` is concatenation.
Parameters
----------
gnn_module : nn.Module
GNN module for message passing. :attr:`GroupRevRes` will clone the module for
:attr:`groups`-1 number of times, yielding :attr:`groups` copies in total.
The input and output node representation size need to be the same. Its forward
function needs to take a DGLGraph and the associated input node features in order,
optionally followed by additional arguments like edge features.
groups : int, optional
The number of groups.
Examples
--------
>>> import dgl
>>> import torch
>>> import torch.nn as nn
>>> from dgl.nn import GraphConv, GroupRevRes
>>> class GNNLayer(nn.Module):
... def __init__(self, feats, dropout=0.2):
... super(GNNLayer, self).__init__()
... # Use BatchNorm and dropout to prevent gradient vanishing
... # In particular if you use a large number of GNN layers
... self.norm = nn.BatchNorm1d(feats)
... self.conv = GraphConv(feats, feats)
... self.dropout = nn.Dropout(dropout)
...
... def forward(self, g, x):
... x = self.norm(x)
... x = self.dropout(x)
... return self.conv(g, x)
>>> num_nodes = 5
>>> num_edges = 20
>>> feats = 32
>>> groups = 2
>>> g = dgl.rand_graph(num_nodes, num_edges)
>>> x = torch.randn(num_nodes, feats)
>>> conv = GNNLayer(feats // groups)
>>> model = GroupRevRes(conv, groups)
>>> out = model(g, x)
"""
def __init__(self, gnn_module, groups=2):
super(GroupRevRes, self).__init__()
self.gnn_modules = nn.ModuleList()
for i in range(groups):
if i == 0:
self.gnn_modules.append(gnn_module)
else:
self.gnn_modules.append(deepcopy(gnn_module))
self.groups = groups
def _forward(self, g, x, *args):
xs = torch.chunk(x, self.groups, dim=-1)
if len(args) == 0:
args_chunks = [()] * self.groups
else:
chunked_args = list(
map(lambda arg: torch.chunk(arg, self.groups, dim=-1), args)
)
args_chunks = list(zip(*chunked_args))
y_in = sum(xs[1:])
ys = []
for i in range(self.groups):
y_in = xs[i] + self.gnn_modules[i](g, y_in, *args_chunks[i])
ys.append(y_in)
out = torch.cat(ys, dim=-1)
return out
def _inverse(self, g, y, *args):
ys = torch.chunk(y, self.groups, dim=-1)
if len(args) == 0:
args_chunks = [()] * self.groups
else:
chunked_args = list(
map(lambda arg: torch.chunk(arg, self.groups, dim=-1), args)
)
args_chunks = list(zip(*chunked_args))
xs = []
for i in range(self.groups - 1, -1, -1):
if i != 0:
y_in = ys[i - 1]
else:
y_in = sum(xs)
x = ys[i] - self.gnn_modules[i](g, y_in, *args_chunks[i])
xs.append(x)
x = torch.cat(xs[::-1], dim=-1)
return x
def forward(self, g, x, *args):
r"""Apply the GNN module with grouped reversible residual connection.
Parameters
----------
g : DGLGraph
The graph.
x : torch.Tensor
The input feature of shape :math:`(N, D_{in})`, where :math:`D_{in}` is size
of input feature, :math:`N` is the number of nodes.
args
Additional arguments to pass to :attr:`gnn_module`.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, D_{in})`.
"""
args = (g, x) + args
y = InvertibleCheckpoint.apply(
self._forward,
self._inverse,
len(args),
*(args + tuple([p for p in self.parameters() if p.requires_grad]))
)
return y
python/dgl/nn/pytorch/conv/hgtconv.py
+201
View File
@@ -0,0 +1,201 @@
"""Heterogeneous Graph Transformer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import math
import torch
import torch.nn as nn
from .... import function as fn
from ..linear import TypedLinear
from ..softmax import edge_softmax
class HGTConv(nn.Module):
r"""Heterogeneous graph transformer convolution from `Heterogeneous Graph Transformer
<https://arxiv.org/abs/2003.01332>`__
Given a graph :math:`G(V, E)` and input node features :math:`H^{(l-1)}`,
it computes the new node features as follows:
Compute a multi-head attention score for each edge :math:`(s, e, t)` in the graph:
.. math::
Attention(s, e, t) = \text{Softmax}\left(||_{i\in[1,h]}ATT-head^i(s, e, t)\right) \\
ATT-head^i(s, e, t) = \left(K^i(s)W^{ATT}_{\phi(e)}Q^i(t)^{\top}\right)\cdot
\frac{\mu_{(\tau(s),\phi(e),\tau(t)}}{\sqrt{d}} \\
K^i(s) = \text{K-Linear}^i_{\tau(s)}(H^{(l-1)}[s]) \\
Q^i(t) = \text{Q-Linear}^i_{\tau(t)}(H^{(l-1)}[t]) \\
Compute the message to send on each edge :math:`(s, e, t)`:
.. math::
Message(s, e, t) = ||_{i\in[1, h]} MSG-head^i(s, e, t) \\
MSG-head^i(s, e, t) = \text{M-Linear}^i_{\tau(s)}(H^{(l-1)}[s])W^{MSG}_{\phi(e)} \\
Send messages to target nodes :math:`t` and aggregate:
.. math::
\tilde{H}^{(l)}[t] = \sum_{\forall s\in \mathcal{N}(t)}\left( Attention(s,e,t)
\cdot Message(s,e,t)\right)
Compute new node features:
.. math::
H^{(l)}[t]=\text{A-Linear}_{\tau(t)}(\sigma(\tilde(H)^{(l)}[t])) + H^{(l-1)}[t]
Parameters
----------
in_size : int
Input node feature size.
head_size : int
Output head size. The output node feature size is ``head_size * num_heads``.
num_heads : int
Number of heads. The output node feature size is ``head_size * num_heads``.
num_ntypes : int
Number of node types.
num_etypes : int
Number of edge types.
dropout : optional, float
Dropout rate.
use_norm : optiona, bool
If true, apply a layer norm on the output node feature.
Examples
--------
"""
def __init__(
self,
in_size,
head_size,
num_heads,
num_ntypes,
num_etypes,
dropout=0.2,
use_norm=False,
):
super().__init__()
self.in_size = in_size
self.head_size = head_size
self.num_heads = num_heads
self.sqrt_d = math.sqrt(head_size)
self.use_norm = use_norm
self.linear_k = TypedLinear(in_size, head_size * num_heads, num_ntypes)
self.linear_q = TypedLinear(in_size, head_size * num_heads, num_ntypes)
self.linear_v = TypedLinear(in_size, head_size * num_heads, num_ntypes)
self.linear_a = TypedLinear(
head_size * num_heads, head_size * num_heads, num_ntypes
)
self.relation_pri = nn.ParameterList(
[nn.Parameter(torch.ones(num_etypes)) for i in range(num_heads)]
)
self.relation_att = nn.ModuleList(
[
TypedLinear(head_size, head_size, num_etypes)
for i in range(num_heads)
]
)
self.relation_msg = nn.ModuleList(
[
TypedLinear(head_size, head_size, num_etypes)
for i in range(num_heads)
]
)
self.skip = nn.Parameter(torch.ones(num_ntypes))
self.drop = nn.Dropout(dropout)
if use_norm:
self.norm = nn.LayerNorm(head_size * num_heads)
if in_size != head_size * num_heads:
self.residual_w = nn.Parameter(
torch.Tensor(in_size, head_size * num_heads)
)
nn.init.xavier_uniform_(self.residual_w)
def forward(self, g, x, ntype, etype, *, presorted=False):
"""Forward computation.
Parameters
----------
g : DGLGraph
The input graph.
x : torch.Tensor
A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
ntype : torch.Tensor
An 1D integer tensor of node types. Shape: :math:`(|V|,)`.
etype : torch.Tensor
An 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
presorted : bool, optional
Whether *both* the nodes and the edges of the input graph have been sorted by
their types. Forward on pre-sorted graph may be faster. Graphs created by
:func:`~dgl.to_homogeneous` automatically satisfy the condition.
Also see :func:`~dgl.reorder_graph` for manually reordering the nodes and edges.
Returns
-------
torch.Tensor
New node features. Shape: :math:`(|V|, D_{head} * N_{head})`.
"""
self.presorted = presorted
if g.is_block:
x_src = x
x_dst = x[: g.num_dst_nodes()]
srcntype = ntype
dstntype = ntype[: g.num_dst_nodes()]
else:
x_src = x
x_dst = x
srcntype = ntype
dstntype = ntype
with g.local_scope():
k = self.linear_k(x_src, srcntype, presorted).view(
-1, self.num_heads, self.head_size
)
q = self.linear_q(x_dst, dstntype, presorted).view(
-1, self.num_heads, self.head_size
)
v = self.linear_v(x_src, srcntype, presorted).view(
-1, self.num_heads, self.head_size
)
g.srcdata["k"] = k
g.dstdata["q"] = q
g.srcdata["v"] = v
g.edata["etype"] = etype
g.apply_edges(self.message)
g.edata["m"] = g.edata["m"] * edge_softmax(
g, g.edata["a"]
).unsqueeze(-1)
g.update_all(fn.copy_e("m", "m"), fn.sum("m", "h"))
h = g.dstdata["h"].view(-1, self.num_heads * self.head_size)
# target-specific aggregation
h = self.drop(self.linear_a(h, dstntype, presorted))
alpha = torch.sigmoid(self.skip[dstntype]).unsqueeze(-1)
if x_dst.shape != h.shape:
h = h * alpha + (x_dst @ self.residual_w) * (1 - alpha)
else:
h = h * alpha + x_dst * (1 - alpha)
if self.use_norm:
h = self.norm(h)
return h
def message(self, edges):
"""Message function."""
a, m = [], []
etype = edges.data["etype"]
k = torch.unbind(edges.src["k"], dim=1)
q = torch.unbind(edges.dst["q"], dim=1)
v = torch.unbind(edges.src["v"], dim=1)
for i in range(self.num_heads):
kw = self.relation_att[i](k[i], etype, self.presorted) # (E, O)
a.append(
(kw * q[i]).sum(-1) * self.relation_pri[i][etype] / self.sqrt_d
) # (E,)
m.append(
self.relation_msg[i](v[i], etype, self.presorted)
) # (E, O)
return {"a": torch.stack(a, dim=1), "m": torch.stack(m, dim=1)}
python/dgl/nn/pytorch/conv/nnconv.py
+187
View File
@@ -0,0 +1,187 @@
"""Torch Module for NNConv layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from torch.nn import init
from .... import function as fn
from ....utils import expand_as_pair
from ..utils import Identity
class NNConv(nn.Module):
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)}`.
NNConv 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 torch as th
>>> 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 = th.ones(6, 10)
>>> lin = th.nn.Linear(5, 20)
>>> def edge_func(efeat):
... return lin(efeat)
>>> efeat = th.ones(6+6, 5)
>>> conv = NNConv(10, 2, edge_func, 'mean')
>>> res = conv(g, feat, efeat)
>>> res
tensor([[-1.5243, -0.2719],
[-1.5243, -0.2719],
[-1.5243, -0.2719],
[-1.5243, -0.2719],
[-1.5243, -0.2719],
[-1.5243, -0.2719]], grad_fn=<AddBackward0>)
>>> # 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 = th.tensor(np.random.rand(2, 10).astype(np.float32))
>>> v_feat = th.tensor(np.random.rand(4, 10).astype(np.float32))
>>> conv = NNConv(10, 2, edge_func, 'mean')
>>> efeat = th.ones(5, 5)
>>> res = conv(g, (u_feat, v_feat), efeat)
>>> res
tensor([[-0.6568, 0.5042],
[ 0.9089, -0.5352],
[ 0.1261, -0.0155],
[-0.6568, 0.5042]], grad_fn=<AddBackward0>)
"""
def __init__(
self,
in_feats,
out_feats,
edge_func,
aggregator_type="mean",
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
self.edge_func = edge_func
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
if residual:
if self._in_dst_feats != out_feats:
self.res_fc = nn.Linear(
self._in_dst_feats, out_feats, bias=False
)
else:
self.res_fc = Identity()
else:
self.register_buffer("res_fc", None)
if bias:
self.bias = nn.Parameter(th.Tensor(out_feats))
else:
self.register_buffer("bias", None)
self.reset_parameters()
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The model parameters are initialized using Glorot uniform initialization
and the bias is initialized to be zero.
"""
gain = init.calculate_gain("relu")
if self.bias is not None:
nn.init.zeros_(self.bias)
if isinstance(self.res_fc, nn.Linear):
nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
def forward(self, graph, feat, efeat):
r"""Compute MPNN Graph Convolution layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor or pair of torch.Tensor
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 : torch.Tensor
The edge feature of shape :math:`(E, *)`, which should fit the input
shape requirement of ``edge_func``. :math:`E` is the number of edges
of the graph.
Returns
-------
torch.Tensor
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.unsqueeze(-1)
# (n, d_in, d_out)
graph.edata["w"] = self.edge_func(efeat).view(
-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["neigh"].sum(dim=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
return rst
python/dgl/nn/pytorch/conv/pnaconv.py
+347
View File
@@ -0,0 +1,347 @@
"""Torch Module for Principal Neighbourhood Aggregation Convolution Layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import numpy as np
import torch
import torch.nn as nn
def aggregate_mean(h):
"""mean aggregation"""
return torch.mean(h, dim=1)
def aggregate_max(h):
"""max aggregation"""
return torch.max(h, dim=1)[0]
def aggregate_min(h):
"""min aggregation"""
return torch.min(h, dim=1)[0]
def aggregate_sum(h):
"""sum aggregation"""
return torch.sum(h, dim=1)
def aggregate_std(h):
"""standard deviation aggregation"""
return torch.sqrt(aggregate_var(h) + 1e-30)
def aggregate_var(h):
"""variance aggregation"""
h_mean_squares = torch.mean(h * h, dim=1)
h_mean = torch.mean(h, dim=1)
var = torch.relu(h_mean_squares - h_mean * h_mean)
return var
def _aggregate_moment(h, n):
"""moment aggregation: for each node (E[(X-E[X])^n])^{1/n}"""
h_mean = torch.mean(h, dim=1, keepdim=True)
h_n = torch.mean(torch.pow(h - h_mean, n), dim=1)
rooted_h_n = torch.sign(h_n) * torch.pow(torch.abs(h_n) + 1e-30, 1.0 / n)
return rooted_h_n
def aggregate_moment_3(h):
"""moment aggregation with n=3"""
return _aggregate_moment(h, n=3)
def aggregate_moment_4(h):
"""moment aggregation with n=4"""
return _aggregate_moment(h, n=4)
def aggregate_moment_5(h):
"""moment aggregation with n=5"""
return _aggregate_moment(h, n=5)
def scale_identity(h):
"""identity scaling (no scaling operation)"""
return h
def scale_amplification(h, D, delta):
"""amplification scaling"""
return h * (np.log(D + 1) / delta)
def scale_attenuation(h, D, delta):
"""attenuation scaling"""
return h * (delta / np.log(D + 1))
AGGREGATORS = {
"mean": aggregate_mean,
"sum": aggregate_sum,
"max": aggregate_max,
"min": aggregate_min,
"std": aggregate_std,
"var": aggregate_var,
"moment3": aggregate_moment_3,
"moment4": aggregate_moment_4,
"moment5": aggregate_moment_5,
}
SCALERS = {
"identity": scale_identity,
"amplification": scale_amplification,
"attenuation": scale_attenuation,
}
class PNAConvTower(nn.Module):
"""A single PNA tower in PNA layers"""
def __init__(
self,
in_size,
out_size,
aggregators,
scalers,
delta,
dropout=0.0,
edge_feat_size=0,
):
super(PNAConvTower, self).__init__()
self.in_size = in_size
self.out_size = out_size
self.aggregators = aggregators
self.scalers = scalers
self.delta = delta
self.edge_feat_size = edge_feat_size
self.M = nn.Linear(2 * in_size + edge_feat_size, in_size)
self.U = nn.Linear(
(len(aggregators) * len(scalers) + 1) * in_size, out_size
)
self.dropout = nn.Dropout(dropout)
self.batchnorm = nn.BatchNorm1d(out_size)
def reduce_func(self, nodes):
"""reduce function for PNA layer:
tensordot of multiple aggregation and scaling operations"""
msg = nodes.mailbox["msg"]
degree = msg.size(1)
h = torch.cat(
[AGGREGATORS[agg](msg) for agg in self.aggregators], dim=1
)
h = torch.cat(
[
SCALERS[scaler](h, D=degree, delta=self.delta)
if scaler != "identity"
else h
for scaler in self.scalers
],
dim=1,
)
return {"h_neigh": h}
def message(self, edges):
"""message function for PNA layer"""
if self.edge_feat_size > 0:
f = torch.cat(
[edges.src["h"], edges.dst["h"], edges.data["a"]], dim=-1
)
else:
f = torch.cat([edges.src["h"], edges.dst["h"]], dim=-1)
return {"msg": self.M(f)}
def forward(self, graph, node_feat, edge_feat=None):
"""compute the forward pass of a single tower in PNA convolution layer"""
# calculate graph normalization factors
snorm_n = torch.cat(
[
torch.ones(N, 1).to(node_feat) / N
for N in graph.batch_num_nodes()
],
dim=0,
).sqrt()
with graph.local_scope():
graph.ndata["h"] = node_feat
if self.edge_feat_size > 0:
assert edge_feat is not None, "Edge features must be provided."
graph.edata["a"] = edge_feat
graph.update_all(self.message, self.reduce_func)
h = self.U(torch.cat([node_feat, graph.ndata["h_neigh"]], dim=-1))
h = h * snorm_n
return self.dropout(self.batchnorm(h))
class PNAConv(nn.Module):
r"""Principal Neighbourhood Aggregation Layer from `Principal Neighbourhood Aggregation
for Graph Nets <https://arxiv.org/abs/2004.05718>`__
A PNA layer is composed of multiple PNA towers. Each tower takes as input a split of the
input features, and computes the message passing as below.
.. math::
h_i^(l+1) = U(h_i^l, \oplus_{(i,j)\in E}M(h_i^l, e_{i,j}, h_j^l))
where :math:`h_i` and :math:`e_{i,j}` are node features and edge features, respectively.
:math:`M` and :math:`U` are MLPs, taking the concatenation of input for computing
output features. :math:`\oplus` represents the combination of various aggregators
and scalers. Aggregators aggregate messages from neighbours and scalers scale the
aggregated messages in different ways. :math:`\oplus` concatenates the output features
of each combination.
The output of multiple towers are concatenated and fed into a linear mixing layer for the
final output.
Parameters
----------
in_size : int
Input feature size; i.e. the size of :math:`h_i^l`.
out_size : int
Output feature size; i.e. the size of :math:`h_i^{l+1}`.
aggregators : list of str
List of aggregation function names(each aggregator specifies a way to aggregate
messages from neighbours), selected from:
* ``mean``: the mean of neighbour messages
* ``max``: the maximum of neighbour messages
* ``min``: the minimum of neighbour messages
* ``std``: the standard deviation of neighbour messages
* ``var``: the variance of neighbour messages
* ``sum``: the sum of neighbour messages
* ``moment3``, ``moment4``, ``moment5``: the normalized moments aggregation
:math:`(E[(X-E[X])^n])^{1/n}`
scalers: list of str
List of scaler function names, selected from:
* ``identity``: no scaling
* ``amplification``: multiply the aggregated message by :math:`\log(d+1)/\delta`,
where :math:`d` is the degree of the node.
* ``attenuation``: multiply the aggregated message by :math:`\delta/\log(d+1)`
delta: float
The degree-related normalization factor computed over the training set, used by scalers
for normalization. :math:`E[\log(d+1)]`, where :math:`d` is the degree for each node
in the training set.
dropout: float, optional
The dropout ratio. Default: 0.0.
num_towers: int, optional
The number of towers used. Default: 1. Note that in_size and out_size must be divisible
by num_towers.
edge_feat_size: int, optional
The edge feature size. Default: 0.
residual : bool, optional
The bool flag that determines whether to add a residual connection for the
output. Default: True. If in_size and out_size of the PNA conv layer are not
the same, this flag will be set as False forcibly.
Example
-------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import PNAConv
>>>
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 10)
>>> conv = PNAConv(10, 10, ['mean', 'max', 'sum'], ['identity', 'amplification'], 2.5)
>>> ret = conv(g, feat)
"""
def __init__(
self,
in_size,
out_size,
aggregators,
scalers,
delta,
dropout=0.0,
num_towers=1,
edge_feat_size=0,
residual=True,
):
super(PNAConv, self).__init__()
self.in_size = in_size
self.out_size = out_size
assert (
in_size % num_towers == 0
), "in_size must be divisible by num_towers"
assert (
out_size % num_towers == 0
), "out_size must be divisible by num_towers"
self.tower_in_size = in_size // num_towers
self.tower_out_size = out_size // num_towers
self.edge_feat_size = edge_feat_size
self.residual = residual
if self.in_size != self.out_size:
self.residual = False
self.towers = nn.ModuleList(
[
PNAConvTower(
self.tower_in_size,
self.tower_out_size,
aggregators,
scalers,
delta,
dropout=dropout,
edge_feat_size=edge_feat_size,
)
for _ in range(num_towers)
]
)
self.mixing_layer = nn.Sequential(
nn.Linear(out_size, out_size), nn.LeakyReLU()
)
def forward(self, graph, node_feat, edge_feat=None):
r"""
Description
-----------
Compute PNA layer.
Parameters
----------
graph : DGLGraph
The graph.
node_feat : torch.Tensor
The input feature of shape :math:`(N, h_n)`. :math:`N` is the number of
nodes, and :math:`h_n` must be the same as in_size.
edge_feat : torch.Tensor, optional
The edge feature of shape :math:`(M, h_e)`. :math:`M` is the number of
edges, and :math:`h_e` must be the same as edge_feat_size.
Returns
-------
torch.Tensor
The output node feature of shape :math:`(N, h_n')` where :math:`h_n'`
should be the same as out_size.
"""
h_cat = torch.cat(
[
tower(
graph,
node_feat[
:,
ti * self.tower_in_size : (ti + 1) * self.tower_in_size,
],
edge_feat,
)
for ti, tower in enumerate(self.towers)
],
dim=1,
)
h_out = self.mixing_layer(h_cat)
# add residual connection
if self.residual:
h_out = h_out + node_feat
return h_out
python/dgl/nn/pytorch/conv/relgraphconv.py
+195
View File
@@ -0,0 +1,195 @@
"""Torch Module for Relational graph convolution layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from .... import function as fn
from ..linear import TypedLinear
class RelGraphConv(nn.Module):
r"""Relational graph convolution layer from `Modeling Relational Data with Graph
Convolutional Networks <https://arxiv.org/abs/1703.06103>`__
It can be described in as below:
.. math::
h_i^{(l+1)} = \sigma(\sum_{r\in\mathcal{R}}
\sum_{j\in\mathcal{N}^r(i)}e_{j,i}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:`e_{j,i}` is the normalizer. :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, optional
Which weight regularizer to use ("basis", "bdd" or ``None``):
- "basis" is for basis-decomposition.
- "bdd" is for block-diagonal-decomposition.
- ``None`` applies no regularization.
Default: ``None``.
num_bases : int, optional
Number of bases. It comes into effect when a regularizer is applied.
If ``None``, it uses number of relations (``num_rels``). Default: ``None``.
Note that ``in_feat`` and ``out_feat`` must be divisible by ``num_bases``
when applying "bdd" regularizer.
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``.
dropout : float, optional
Dropout rate. Default: ``0.0``
layer_norm: bool, optional
True to add layer norm. Default: ``False``
Examples
--------
>>> import dgl
>>> import numpy as np
>>> import torch as th
>>> from dgl.nn import RelGraphConv
>>>
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 10)
>>> conv = RelGraphConv(10, 2, 3, regularizer='basis', num_bases=2)
>>> etype = th.tensor([0,1,2,0,1,2])
>>> res = conv(g, feat, etype)
>>> res
tensor([[ 0.3996, -2.3303],
[-0.4323, -0.1440],
[ 0.3996, -2.3303],
[ 2.1046, -2.8654],
[-0.4323, -0.1440],
[-0.1309, -1.0000]], grad_fn=<AddBackward0>)
"""
def __init__(
self,
in_feat,
out_feat,
num_rels,
regularizer=None,
num_bases=None,
bias=True,
activation=None,
self_loop=True,
dropout=0.0,
layer_norm=False,
):
super().__init__()
if regularizer is not None and num_bases is None:
num_bases = num_rels
self.linear_r = TypedLinear(
in_feat, out_feat, num_rels, regularizer, num_bases
)
self.bias = bias
self.activation = activation
self.self_loop = self_loop
self.layer_norm = layer_norm
# bias
if self.bias:
self.h_bias = nn.Parameter(th.Tensor(out_feat))
nn.init.zeros_(self.h_bias)
# TODO(minjie): consider remove those options in the future to make
# the module only about graph convolution.
# layer norm
if self.layer_norm:
self.layer_norm_weight = nn.LayerNorm(
out_feat, elementwise_affine=True
)
# weight for self loop
if self.self_loop:
self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
nn.init.xavier_uniform_(
self.loop_weight, gain=nn.init.calculate_gain("relu")
)
self.dropout = nn.Dropout(dropout)
def message(self, edges):
"""Message function."""
m = self.linear_r(edges.src["h"], edges.data["etype"], self.presorted)
if "norm" in edges.data:
m = m * edges.data["norm"]
return {"m": m}
def forward(self, g, feat, etypes, norm=None, *, presorted=False):
"""Forward computation.
Parameters
----------
g : DGLGraph
The graph.
feat : torch.Tensor
A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
etypes : torch.Tensor or list[int]
An 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
norm : torch.Tensor, optional
An 1D tensor of edge norm value. Shape: :math:`(|E|,)`.
presorted : bool, optional
Whether the edges of the input graph have been sorted by their types.
Forward on pre-sorted graph may be faster. Graphs created
by :func:`~dgl.to_homogeneous` automatically satisfy the condition.
Also see :func:`~dgl.reorder_graph` for sorting edges manually.
Returns
-------
torch.Tensor
New node features. Shape: :math:`(|V|, D_{out})`.
"""
self.presorted = presorted
with g.local_scope():
g.srcdata["h"] = feat
if norm is not None:
g.edata["norm"] = norm
g.edata["etype"] = etypes
# message passing
g.update_all(self.message, fn.sum("m", "h"))
# apply bias and activation
h = g.dstdata["h"]
if self.layer_norm:
h = self.layer_norm_weight(h)
if self.bias:
h = h + self.h_bias
if self.self_loop:
h = h + feat[: g.num_dst_nodes()] @ self.loop_weight
if self.activation:
h = self.activation(h)
h = self.dropout(h)
return h
python/dgl/nn/pytorch/conv/sageconv.py
+295
View File
@@ -0,0 +1,295 @@
"""Torch Module for GraphSAGE layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch
from torch import nn
from torch.nn import functional as F
from .... import function as fn
from ....base import DGLError
from ....utils import check_eq_shape, expand_as_pair
class SAGEConv(nn.Module):
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)})
If a weight tensor on each edge is provided, the aggregation becomes:
.. math::
h_{\mathcal{N}(i)}^{(l+1)} = \mathrm{aggregate}
\left(\{e_{ji} h_{j}^{l}, \forall j \in \mathcal{N}(i) \}\right)
where :math:`e_{ji}` is the scalar weight on the edge from node :math:`j` to node :math:`i`.
Please make sure that :math:`e_{ji}` is broadcastable with :math:`h_j^{l}`.
Parameters
----------
in_feats : int, or pair of ints
Input feature size; i.e, the number of dimensions of :math:`h_i^{(l)}`.
SAGEConv 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 torch as th
>>> 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 = th.ones(6, 10)
>>> conv = SAGEConv(10, 2, 'pool')
>>> res = conv(g, feat)
>>> res
tensor([[-1.0888, -2.1099],
[-1.0888, -2.1099],
[-1.0888, -2.1099],
[-1.0888, -2.1099],
[-1.0888, -2.1099],
[-1.0888, -2.1099]], grad_fn=<AddBackward0>)
>>> # 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 = th.rand(2, 5)
>>> v_fea = th.rand(4, 10)
>>> conv = SAGEConv((5, 10), 2, 'mean')
>>> res = conv(g, (u_fea, v_fea))
>>> res
tensor([[ 0.3163, 3.1166],
[ 0.3866, 2.5398],
[ 0.5873, 1.6597],
[-0.2502, 2.8068]], grad_fn=<AddBackward0>)
"""
def __init__(
self,
in_feats,
out_feats,
aggregator_type,
feat_drop=0.0,
bias=True,
norm=None,
activation=None,
):
super(SAGEConv, self).__init__()
valid_aggre_types = {"mean", "gcn", "pool", "lstm"}
if aggregator_type not in valid_aggre_types:
raise DGLError(
"Invalid aggregator_type. Must be one of {}. "
"But got {!r} instead.".format(
valid_aggre_types, aggregator_type
)
)
self._in_src_feats, self._in_dst_feats = expand_as_pair(in_feats)
self._out_feats = out_feats
self._aggre_type = aggregator_type
self.norm = norm
self.feat_drop = nn.Dropout(feat_drop)
self.activation = activation
# aggregator type: mean/pool/lstm/gcn
if aggregator_type == "pool":
self.fc_pool = nn.Linear(self._in_src_feats, self._in_src_feats)
if aggregator_type == "lstm":
self.lstm = nn.LSTM(
self._in_src_feats, self._in_src_feats, batch_first=True
)
self.fc_neigh = nn.Linear(self._in_src_feats, out_feats, bias=False)
if aggregator_type != "gcn":
self.fc_self = nn.Linear(self._in_dst_feats, out_feats, bias=bias)
elif bias:
self.bias = nn.parameter.Parameter(torch.zeros(self._out_feats))
else:
self.register_buffer("bias", None)
self.reset_parameters()
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The linear weights :math:`W^{(l)}` are initialized using Glorot uniform initialization.
The LSTM module is using xavier initialization method for its weights.
"""
gain = nn.init.calculate_gain("relu")
if self._aggre_type == "pool":
nn.init.xavier_uniform_(self.fc_pool.weight, gain=gain)
if self._aggre_type == "lstm":
self.lstm.reset_parameters()
if self._aggre_type != "gcn":
nn.init.xavier_uniform_(self.fc_self.weight, gain=gain)
nn.init.xavier_uniform_(self.fc_neigh.weight, gain=gain)
def _lstm_reducer(self, nodes):
"""LSTM reducer
NOTE(zihao): lstm reducer with default schedule (degree bucketing)
is slow, we could accelerate this with degree padding in the future.
"""
m = nodes.mailbox["m"] # (B, L, D)
batch_size = m.shape[0]
h = (
m.new_zeros((1, batch_size, self._in_src_feats)),
m.new_zeros((1, batch_size, self._in_src_feats)),
)
_, (rst, _) = self.lstm(m, h)
return {"neigh": rst.squeeze(0)}
def forward(self, graph, feat, edge_weight=None):
r"""
Description
-----------
Compute GraphSAGE layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor or pair of torch.Tensor
If a torch.Tensor is given, it represents the input feature of shape
:math:`(N, D_{in})`
where :math:`D_{in}` is size of input feature, :math:`N` is the number of nodes.
If a pair of torch.Tensor is given, the pair must contain two tensors of shape
:math:`(N_{in}, D_{in_{src}})` and :math:`(N_{out}, D_{in_{dst}})`.
edge_weight : torch.Tensor, optional
Optional tensor on the edge. If given, the convolution will weight
with regard to the message.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N_{dst}, D_{out})`
where :math:`N_{dst}` is the number of destination nodes in the input graph,
:math:`D_{out}` is the size of the 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()]
msg_fn = fn.copy_u("h", "m")
if edge_weight is not None:
assert edge_weight.shape[0] == graph.num_edges()
graph.edata["_edge_weight"] = edge_weight
msg_fn = fn.u_mul_e("h", "_edge_weight", "m")
h_self = feat_dst
# Handle the case of graphs without edges
if graph.num_edges() == 0:
graph.dstdata["neigh"] = torch.zeros(
feat_dst.shape[0], self._in_src_feats
).to(feat_dst)
# Determine whether to apply linear transformation before message passing A(XW)
lin_before_mp = self._in_src_feats > self._out_feats
# Message Passing
if self._aggre_type == "mean":
graph.srcdata["h"] = (
self.fc_neigh(feat_src) if lin_before_mp else feat_src
)
graph.update_all(msg_fn, fn.mean("m", "neigh"))
h_neigh = graph.dstdata["neigh"]
if not lin_before_mp:
h_neigh = self.fc_neigh(h_neigh)
elif self._aggre_type == "gcn":
check_eq_shape(feat)
graph.srcdata["h"] = (
self.fc_neigh(feat_src) if lin_before_mp else feat_src
)
if isinstance(feat, tuple): # heterogeneous
graph.dstdata["h"] = (
self.fc_neigh(feat_dst) if lin_before_mp else feat_dst
)
else:
if graph.is_block:
graph.dstdata["h"] = graph.srcdata["h"][
: graph.num_dst_nodes()
]
else:
graph.dstdata["h"] = graph.srcdata["h"]
graph.update_all(msg_fn, fn.sum("m", "neigh"))
# divide in_degrees
degs = graph.in_degrees().to(feat_dst)
h_neigh = (graph.dstdata["neigh"] + graph.dstdata["h"]) / (
degs.unsqueeze(-1) + 1
)
if not lin_before_mp:
h_neigh = self.fc_neigh(h_neigh)
elif self._aggre_type == "pool":
graph.srcdata["h"] = F.relu(self.fc_pool(feat_src))
graph.update_all(msg_fn, fn.max("m", "neigh"))
h_neigh = self.fc_neigh(graph.dstdata["neigh"])
elif self._aggre_type == "lstm":
graph.srcdata["h"] = feat_src
graph.update_all(msg_fn, self._lstm_reducer)
h_neigh = self.fc_neigh(graph.dstdata["neigh"])
else:
raise KeyError(
"Aggregator type {} not recognized.".format(
self._aggre_type
)
)
# GraphSAGE GCN does not require fc_self.
if self._aggre_type == "gcn":
rst = h_neigh
# add bias manually for GCN
if self.bias is not None:
rst = rst + self.bias
else:
rst = self.fc_self(h_self) + h_neigh
# activation
if self.activation is not None:
rst = self.activation(rst)
# normalization
if self.norm is not None:
rst = self.norm(rst)
return rst
python/dgl/nn/pytorch/conv/sgconv.py
+218
View File
@@ -0,0 +1,218 @@
"""Torch Module for Simplifying Graph Convolution layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from .... import function as fn
from ....base import DGLError
from .graphconv import EdgeWeightNorm
class SGConv(nn.Module):
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 torch as th
>>> 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 = th.ones(6, 10)
>>> conv = SGConv(10, 2, k=2)
>>> res = conv(g, feat)
>>> res
tensor([[-1.9441, -0.9343],
[-1.9441, -0.9343],
[-1.9441, -0.9343],
[-2.7709, -1.3316],
[-1.9297, -0.9273],
[-1.9441, -0.9343]], grad_fn=<AddmmBackward>)
"""
def __init__(
self,
in_feats,
out_feats,
k=1,
cached=False,
bias=True,
norm=None,
allow_zero_in_degree=False,
):
super(SGConv, self).__init__()
self.fc = nn.Linear(in_feats, out_feats, bias=bias)
self._cached = cached
self._cached_h = None
self._k = k
self.norm = norm
self._allow_zero_in_degree = allow_zero_in_degree
self.reset_parameters()
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The model parameters are initialized using xavier initialization
and the bias is initialized to be zero.
"""
nn.init.xavier_uniform_(self.fc.weight)
if self.fc.bias is not None:
nn.init.zeros_(self.fc.bias)
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, edge_weight=None):
r"""
Description
-----------
Compute Simplifying Graph Convolution layer.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor
The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
is size of input feature, :math:`N` is the number of nodes.
edge_weight: torch.Tensor, optional
edge_weight to use in the message passing process. This is equivalent to
using weighted adjacency matrix in the equation above, and
:math:`\tilde{D}^{-1/2}\tilde{A} \tilde{D}^{-1/2}`
is based on :class:`dgl.nn.pytorch.conv.graphconv.EdgeWeightNorm`.
Returns
-------
torch.Tensor
The output feature of shape :math:`(N, D_{out})` where :math:`D_{out}`
is size of output feature.
Raises
------
DGLError
If there are 0-in-degree nodes in the input graph, it will raise DGLError
since no message will be passed to those nodes. This will cause invalid output.
The error can be ignored by setting ``allow_zero_in_degree`` parameter to ``True``.
Note
----
If ``cache`` is set to True, ``feat`` and ``graph`` should not change during
training, or you will get wrong results.
"""
with graph.local_scope():
if not self._allow_zero_in_degree:
if (graph.in_degrees() == 0).any():
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."
)
msg_func = fn.copy_u("h", "m")
if edge_weight is not None:
graph.edata["_edge_weight"] = EdgeWeightNorm("both")(
graph, edge_weight
)
msg_func = fn.u_mul_e("h", "_edge_weight", "m")
if self._cached_h is not None:
feat = self._cached_h
else:
if edge_weight is None:
# compute normalization
degs = graph.in_degrees().to(feat).clamp(min=1)
norm = th.pow(degs, -0.5)
norm = norm.to(feat.device).unsqueeze(1)
# compute (D^-1 A^k D)^k X
for _ in range(self._k):
if edge_weight is None:
feat = feat * norm
graph.ndata["h"] = feat
graph.update_all(msg_func, fn.sum("m", "h"))
feat = graph.ndata.pop("h")
if edge_weight is None:
feat = feat * norm
if self.norm is not None:
feat = self.norm(feat)
# cache feature
if self._cached:
self._cached_h = feat
return self.fc(feat)
python/dgl/nn/pytorch/conv/tagconv.py
+150
View File
@@ -0,0 +1,150 @@
"""Torch Module for Topology Adaptive Graph Convolutional layer"""
# pylint: disable= no-member, arguments-differ, invalid-name
import torch as th
from torch import nn
from .... import function as fn
from .graphconv import EdgeWeightNorm
class TAGConv(nn.Module):
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 torch as th
>>> from dgl.nn import TAGConv
>>>
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 10)
>>> conv = TAGConv(10, 2, k=2)
>>> res = conv(g, feat)
>>> res
tensor([[ 0.5490, -1.6373],
[ 0.5490, -1.6373],
[ 0.5490, -1.6373],
[ 0.5513, -1.8208],
[ 0.5215, -1.6044],
[ 0.3304, -1.9927]], grad_fn=<AddmmBackward>)
"""
def __init__(
self,
in_feats,
out_feats,
k=2,
bias=True,
activation=None,
):
super(TAGConv, self).__init__()
self._in_feats = in_feats
self._out_feats = out_feats
self._k = k
self._activation = activation
self.lin = nn.Linear(in_feats * (self._k + 1), out_feats, bias=bias)
self.reset_parameters()
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
Note
----
The model parameters are initialized using Glorot uniform initialization.
"""
gain = nn.init.calculate_gain("relu")
nn.init.xavier_normal_(self.lin.weight, gain=gain)
def forward(self, graph, feat, edge_weight=None):
r"""
Description
-----------
Compute topology adaptive graph convolution.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor
The input feature of shape :math:`(N, D_{in})` where :math:`D_{in}`
is size of input feature, :math:`N` is the number of nodes.
edge_weight: torch.Tensor, optional
edge_weight to use in the message passing process. This is equivalent to
using weighted adjacency matrix in the equation above, and
:math:`\tilde{D}^{-1/2}\tilde{A} \tilde{D}^{-1/2}`
is based on :class:`dgl.nn.pytorch.conv.graphconv.EdgeWeightNorm`.
Returns
-------
torch.Tensor
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"
if edge_weight is None:
norm = th.pow(graph.in_degrees().to(feat).clamp(min=1), -0.5)
shp = norm.shape + (1,) * (feat.dim() - 1)
norm = th.reshape(norm, shp).to(feat.device)
msg_func = fn.copy_u("h", "m")
if edge_weight is not None:
graph.edata["_edge_weight"] = EdgeWeightNorm("both")(
graph, edge_weight
)
msg_func = fn.u_mul_e("h", "_edge_weight", "m")
# D-1/2 A D -1/2 X
fstack = [feat]
for _ in range(self._k):
if edge_weight is None:
rst = fstack[-1] * norm
else:
rst = fstack[-1]
graph.ndata["h"] = rst
graph.update_all(msg_func, fn.sum(msg="m", out="h"))
rst = graph.ndata["h"]
if edge_weight is None:
rst = rst * norm
fstack.append(rst)
rst = self.lin(th.cat(fstack, dim=-1))
if self._activation is not None:
rst = self._activation(rst)
return rst
python/dgl/nn/pytorch/conv/twirlsconv.py
+699
View File
@@ -0,0 +1,699 @@
"""Torch modules for TWIRLS"""
# pylint: disable=invalid-name, useless-super-delegation, no-member
import torch as tc
import torch.nn as nn
import torch.nn.functional as F
from .... import function as fn
class TWIRLSConv(nn.Module):
r"""Convolution together with iteratively reweighting least squre from
`Graph Neural Networks Inspired by Classical Iterative Algorithms
<https://arxiv.org/pdf/2103.06064.pdf>`__
Parameters
----------
input_d : int
Number of input features.
output_d : int
Number of output features.
hidden_d : int
Size of hidden layers.
prop_step : int
Number of propagation steps
num_mlp_before : int
Number of mlp layers before propagation. Default: ``1``.
num_mlp_after : int
Number of mlp layers after propagation. Default: ``1``.
norm : str
The type of norm layers inside mlp layers. Can be ``'batch'``, ``'layer'`` or ``'none'``.
Default: ``'none'``
precond : str
If True, use pre conditioning and unormalized laplacian, else not use pre conditioning
and use normalized laplacian. Default: ``True``
alp : float
The :math:`\alpha` in paper. If equal to :math:`0`, will be automatically decided based
on other hyper prameters. Default: ``0``.
lam : float
The :math:`\lambda` in paper. Default: ``1``.
attention : bool
If ``True``, add an attention layer inside propagations. Default: ``False``.
tau : float
The :math:`\tau` in paper. Default: ``0.2``.
T : float
The :math:`T` in paper. If < 0, :math:`T` will be set to `\infty`. Default: ``-1``.
p : float
The :math:`p` in paper. Default: ``1``.
use_eta : bool
If ``True``, add a learnable weight on each dimension in attention. Default: ``False``.
attn_bef : bool
If ``True``, add another attention layer before propagation. Default: ``False``.
dropout : float
The dropout rate in mlp layers. Default: ``0.0``.
attn_dropout : float
The dropout rate of attention values. Default: ``0.0``.
inp_dropout : float
The dropout rate on input features. Default: ``0.0``.
Note
----
``add_self_loop`` will be automatically called before propagation.
Example
-------
>>> import dgl
>>> from dgl.nn import TWIRLSConv
>>> import torch as th
>>> g = dgl.graph(([0,1,2,3,2,5], [1,2,3,4,0,3]))
>>> feat = th.ones(6, 10)
>>> conv = TWIRLSConv(10, 2, 128, prop_step = 64)
>>> res = conv(g , feat)
>>> res.size()
torch.Size([6, 2])
"""
def __init__(
self,
input_d,
output_d,
hidden_d,
prop_step,
num_mlp_before=1,
num_mlp_after=1,
norm="none",
precond=True,
alp=0,
lam=1,
attention=False,
tau=0.2,
T=-1,
p=1,
use_eta=False,
attn_bef=False,
dropout=0.0,
attn_dropout=0.0,
inp_dropout=0.0,
):
super().__init__()
self.input_d = input_d
self.output_d = output_d
self.hidden_d = hidden_d
self.prop_step = prop_step
self.num_mlp_before = num_mlp_before
self.num_mlp_after = num_mlp_after
self.norm = norm
self.precond = precond
self.attention = attention
self.alp = alp
self.lam = lam
self.tau = tau
self.T = T
self.p = p
self.use_eta = use_eta
self.init_att = attn_bef
self.dropout = dropout
self.attn_dropout = attn_dropout
self.inp_dropout = inp_dropout
# ----- initialization of some variables -----
# where to put attention
self.attn_aft = prop_step // 2 if attention else -1
# whether we can cache unfolding result
self.cacheable = (
(not self.attention)
and self.num_mlp_before == 0
and self.inp_dropout <= 0
)
if self.cacheable:
self.cached_unfolding = None
# if only one layer, then no hidden size
self.size_bef_unf = self.hidden_d
self.size_aft_unf = self.hidden_d
if self.num_mlp_before == 0:
self.size_aft_unf = self.input_d # as the input of mlp_aft
if self.num_mlp_after == 0:
self.size_bef_unf = self.output_d # as the output of mlp_bef
# ----- computational modules -----
self.mlp_bef = MLP(
self.input_d,
self.hidden_d,
self.size_bef_unf,
self.num_mlp_before,
self.dropout,
self.norm,
init_activate=False,
)
self.unfolding = TWIRLSUnfoldingAndAttention(
self.hidden_d,
self.alp,
self.lam,
self.prop_step,
self.attn_aft,
self.tau,
self.T,
self.p,
self.use_eta,
self.init_att,
self.attn_dropout,
self.precond,
)
# if there are really transformations before unfolding, then do init_activate in mlp_aft
self.mlp_aft = MLP(
self.size_aft_unf,
self.hidden_d,
self.output_d,
self.num_mlp_after,
self.dropout,
self.norm,
init_activate=(self.num_mlp_before > 0)
and (self.num_mlp_after > 0),
)
def forward(self, graph, feat):
r"""
Description
-----------
Run TWIRLS forward.
Parameters
----------
graph : DGLGraph
The graph.
feat : torch.Tensor
The initial node features.
Returns
-------
torch.Tensor
The output feature
Note
----
* Input shape: :math:`(N, \text{input_d})` where :math:`N` is the number of nodes.
* Output shape: :math:`(N, \text{output_d})`.
"""
# ensure self loop
graph = graph.remove_self_loop()
graph = graph.add_self_loop()
x = feat
if self.cacheable:
# to cache unfolding result becase there is no paramaters before it
if self.cached_unfolding is None:
self.cached_unfolding = self.unfolding(graph, x)
x = self.cached_unfolding
else:
if self.inp_dropout > 0:
x = F.dropout(x, self.inp_dropout, training=self.training)
x = self.mlp_bef(x)
x = self.unfolding(graph, x)
x = self.mlp_aft(x)
return x
class Propagate(nn.Module):
r"""
Description
-----------
The propagation method which is with pre-conditioning and reparameterizing. Correspond to
eq.28 in the paper.
"""
def __init__(self):
super().__init__()
def _prop(self, graph, Y, lam):
"""propagation part."""
Y = D_power_bias_X(graph, Y, -0.5, lam, 1 - lam)
Y = AX(graph, Y)
Y = D_power_bias_X(graph, Y, -0.5, lam, 1 - lam)
return Y
def forward(self, graph, Y, X, alp, lam):
r"""
Description
-----------
Propagation forward.
Parameters
----------
graph : DGLGraph
The graph.
Y : torch.Tensor
The feature under propagation. Corresponds to :math:`Z^{(k)}` in eq.28 in the paper.
X : torch.Tensor
The original feature. Corresponds to :math:`Z^{(0)}` in eq.28 in the paper.
alp : float
The step size. Corresponds to :math:`\alpha` in the paper.
lam : torch.Tensor
The coefficient of smoothing term. Corresponds to :math:`\lambda` in the paper.
Returns
-------
torch.Tensor
Propagated feature. :math:`Z^{(k+1)}` in eq.28 in the paper.
"""
return (
(1 - alp) * Y
+ alp * lam * self._prop(graph, Y, lam)
+ alp * D_power_bias_X(graph, X, -1, lam, 1 - lam)
)
class PropagateNoPrecond(nn.Module):
r"""
Description
-----------
The propagation method which is without pre-conditioning and reparameterizing and using
normalized laplacian.
Correspond to eq.30 in the paper.
"""
def __init__(self):
super().__init__()
def forward(self, graph, Y, X, alp, lam):
r"""
Description
-----------
Propagation forward.
Parameters
----------
graph : DGLGraph
The graph.
Y : torch.Tensor
The feature under propagation. Corresponds to :math:`Y^{(k)}` in eq.30 in the paper.
X : torch.Tensor
The original feature. Corresponds to :math:`Y^{(0)}` in eq.30 in the paper.
alp : float
The step size. Corresponds to :math:`\alpha` in the paper.
lam : torch.Tensor
The coefficient of smoothing term. Corresponds to :math:`\lambda` in the paper.
Returns
-------
torch.Tensor
Propagated feature. :math:`Y^{(k+1)}` in eq.30 in the paper.
"""
return (
(1 - alp * lam - alp) * Y
+ alp * lam * normalized_AX(graph, Y)
+ alp * X
)
class Attention(nn.Module):
r"""
Description
-----------
The attention function. Correspond to :math:`s` in eq.27 the paper.
Parameters
----------
tau : float
The lower thresholding parameter. Correspond to :math:`\tau` in the paper.
T : float
The upper thresholding parameter. Correspond to :math:`T` in the paper.
p : float
Correspond to :math:`\rho` in the paper..
attn_dropout : float
the dropout rate of attention value. Default: ``0.0``.
Returns
-------
torch.Tensor
The output feature
"""
def __init__(self, tau, T, p, attn_dropout=0.0):
super().__init__()
self.tau = tau
self.T = T
self.p = p
self.attn_dropout = attn_dropout
def reweighting(self, graph):
"""Compute graph edge weight. Would be stored in ``graph.edata['w']``"""
w = graph.edata["w"]
# It is not activation here but to ensure w > 0.
# w can be < 0 here because of some precision issue in dgl, which causes NaN afterwards.
w = F.relu(w) + 1e-7
w = tc.pow(w, 1 - 0.5 * self.p)
w[(w < self.tau)] = self.tau
if self.T > 0:
w[(w > self.T)] = float("inf")
w = 1 / w
# if not (w == w).all():
# raise "nan occured!"
graph.edata["w"] = w + 1e-9 # avoid 0 degree
def forward(self, graph, Y, etas=None):
r"""
Description
-----------
Attention forward. Will update ``graph.edata['w']`` and ``graph.ndata['deg']``.
Parameters
----------
graph : DGLGraph
The graph.
Y : torch.Tensor
The feature to compute attention.
etas : float
The weight of each dimension. If ``None``, then weight of each dimension is 1.
Default: ``None``.
Returns
-------
DGLGraph
The graph.
"""
if etas is not None:
Y = Y * etas.view(-1)
# computing edge distance
graph.srcdata["h"] = Y
graph.srcdata["h_norm"] = (Y**2).sum(-1)
graph.apply_edges(fn.u_dot_v("h", "h", "dot_"))
graph.apply_edges(fn.u_add_v("h_norm", "h_norm", "norm_"))
graph.edata["dot_"] = graph.edata["dot_"].view(-1)
graph.edata["norm_"] = graph.edata["norm_"].view(-1)
graph.edata["w"] = graph.edata["norm_"] - 2 * graph.edata["dot_"]
# apply edge distance to get edge weight
self.reweighting(graph)
# update node degrees
graph.update_all(fn.copy_e("w", "m"), fn.sum("m", "deg"))
graph.ndata["deg"] = graph.ndata["deg"].view(-1)
# attention dropout. the implementation can ensure the degrees do not change in expectation.
# FIXME: consider if there is a better way
if self.attn_dropout > 0:
graph.edata["w"] = F.dropout(
graph.edata["w"], self.attn_dropout, training=self.training
)
return graph
def normalized_AX(graph, X):
"""Y = D^{-1/2}AD^{-1/2}X"""
Y = D_power_X(graph, X, -0.5) # Y = D^{-1/2}X
Y = AX(graph, Y) # Y = AD^{-1/2}X
Y = D_power_X(graph, Y, -0.5) # Y = D^{-1/2}AD^{-1/2}X
return Y
def AX(graph, X):
"""Y = AX"""
graph.srcdata["h"] = X
graph.update_all(
fn.u_mul_e("h", "w", "m"),
fn.sum("m", "h"),
)
Y = graph.dstdata["h"]
return Y
def D_power_X(graph, X, power):
"""Y = D^{power}X"""
degs = graph.ndata["deg"]
norm = tc.pow(degs, power)
Y = X * norm.view(X.size(0), 1)
return Y
def D_power_bias_X(graph, X, power, coeff, bias):
"""Y = (coeff*D + bias*I)^{power} X"""
degs = graph.ndata["deg"]
degs = coeff * degs + bias
norm = tc.pow(degs, power)
Y = X * norm.view(X.size(0), 1)
return Y
class TWIRLSUnfoldingAndAttention(nn.Module):
r"""
Description
-----------
Combine propagation and attention together.
Parameters
----------
d : int
Size of graph feature.
alp : float
Step size. :math:`\alpha` in ther paper.
lam : int
Coefficient of graph smooth term. :math:`\lambda` in ther paper.
prop_step : int
Number of propagation steps
attn_aft : int
Where to put attention layer. i.e. number of propagation steps before attention.
If set to ``-1``, then no attention.
tau : float
The lower thresholding parameter. Correspond to :math:`\tau` in the paper.
T : float
The upper thresholding parameter. Correspond to :math:`T` in the paper.
p : float
Correspond to :math:`\rho` in the paper..
use_eta : bool
If `True`, learn a weight vector for each dimension when doing attention.
init_att : bool
If ``True``, add an extra attention layer before propagation.
attn_dropout : float
the dropout rate of attention value. Default: ``0.0``.
precond : bool
If ``True``, use pre-conditioned & reparameterized version propagation (eq.28), else use
normalized laplacian (eq.30).
Example
-------
>>> import dgl
>>> from dgl.nn import TWIRLSUnfoldingAndAttention
>>> import torch as th
>>> g = dgl.graph(([0, 1, 2, 3, 2, 5], [1, 2, 3, 4, 0, 3])).add_self_loop()
>>> feat = th.ones(6,5)
>>> prop = TWIRLSUnfoldingAndAttention(10, 1, 1, prop_step=3)
>>> res = prop(g,feat)
>>> res
tensor([[2.5000, 2.5000, 2.5000, 2.5000, 2.5000],
[2.5000, 2.5000, 2.5000, 2.5000, 2.5000],
[2.5000, 2.5000, 2.5000, 2.5000, 2.5000],
[3.7656, 3.7656, 3.7656, 3.7656, 3.7656],
[2.5217, 2.5217, 2.5217, 2.5217, 2.5217],
[4.0000, 4.0000, 4.0000, 4.0000, 4.0000]])
"""
def __init__(
self,
d,
alp,
lam,
prop_step,
attn_aft=-1,
tau=0.2,
T=-1,
p=1,
use_eta=False,
init_att=False,
attn_dropout=0,
precond=True,
):
super().__init__()
self.d = d
self.alp = alp if alp > 0 else 1 / (lam + 1) # automatic set alpha
self.lam = lam
self.tau = tau
self.p = p
self.prop_step = prop_step
self.attn_aft = attn_aft
self.use_eta = use_eta
self.init_att = init_att
prop_method = Propagate if precond else PropagateNoPrecond
self.prop_layers = nn.ModuleList(
[prop_method() for _ in range(prop_step)]
)
self.init_attn = (
Attention(tau, T, p, attn_dropout) if self.init_att else None
)
self.attn_layer = (
Attention(tau, T, p, attn_dropout) if self.attn_aft >= 0 else None
)
self.etas = nn.Parameter(tc.ones(d)) if self.use_eta else None
def forward(self, g, X):
r"""
Description
-----------
Compute forward pass of propagation & attention.
Parameters
----------
g : DGLGraph
The graph.
X : torch.Tensor
Init features.
Returns
-------
torch.Tensor
The graph.
"""
Y = X
g.edata["w"] = tc.ones(g.num_edges(), 1, device=g.device)
g.ndata["deg"] = g.in_degrees().to(X)
if self.init_att:
g = self.init_attn(g, Y, self.etas)
for k, layer in enumerate(self.prop_layers):
# do unfolding
Y = layer(g, Y, X, self.alp, self.lam)
# do attention at certain layer
if k == self.attn_aft - 1:
g = self.attn_layer(g, Y, self.etas)
return Y
class MLP(nn.Module):
r"""
Description
-----------
An MLP module.
Parameters
----------
input_d : int
Number of input features.
output_d : int
Number of output features.
hidden_d : int
Size of hidden layers.
num_layers : int
Number of mlp layers.
dropout : float
The dropout rate in mlp layers.
norm : str
The type of norm layers inside mlp layers. Can be ``'batch'``, ``'layer'`` or ``'none'``.
init_activate : bool
If add a relu at the beginning.
"""
def __init__(
self,
input_d,
hidden_d,
output_d,
num_layers,
dropout,
norm,
init_activate,
):
super().__init__()
self.init_activate = init_activate
self.norm = norm
self.dropout = dropout
self.layers = nn.ModuleList([])
if num_layers == 1:
self.layers.append(nn.Linear(input_d, output_d))
elif num_layers > 1:
self.layers.append(nn.Linear(input_d, hidden_d))
for _ in range(num_layers - 2):
self.layers.append(nn.Linear(hidden_d, hidden_d))
self.layers.append(nn.Linear(hidden_d, output_d))
# how many norm layers we have
self.norm_cnt = num_layers - 1 + int(init_activate)
if norm == "batch":
self.norms = nn.ModuleList(
[nn.BatchNorm1d(hidden_d) for _ in range(self.norm_cnt)]
)
elif norm == "layer":
self.norms = nn.ModuleList(
[nn.LayerNorm(hidden_d) for _ in range(self.norm_cnt)]
)
self.reset_params()
def reset_params(self):
"""reset mlp parameters using xavier_norm"""
for layer in self.layers:
nn.init.xavier_normal_(layer.weight.data)
nn.init.constant_(layer.bias.data, 0)
def activate(self, x):
"""do normlaization and activation"""
if self.norm != "none":
x = self.norms[self.cur_norm_idx](x) # use the last norm layer
self.cur_norm_idx += 1
x = F.relu(x)
x = F.dropout(x, self.dropout, training=self.training)
return x
def forward(self, x):
"""The forward pass of mlp."""
self.cur_norm_idx = 0
if self.init_activate:
x = self.activate(x)
for i, layer in enumerate(self.layers):
x = layer(x)
if i != len(self.layers) - 1: # do not activate in the last layer
x = self.activate(x)
return x
python/dgl/nn/pytorch/explain/__init__.py
+6
View File
@@ -0,0 +1,6 @@
"""Torch modules for explanation models."""
# pylint: disable= no-member, arguments-differ, invalid-name
from .gnnexplainer import *
from .subgraphx import *
from .pgexplainer import *
python/dgl/nn/pytorch/explain/gnnexplainer.py
+941
View File
@@ -0,0 +1,941 @@
"""Torch Module for GNNExplainer"""
# pylint: disable= no-member, arguments-differ, invalid-name
from math import sqrt
import torch
from torch import nn
from tqdm.auto import tqdm
from ....base import EID, NID
from ....subgraph import khop_in_subgraph
__all__ = ["GNNExplainer", "HeteroGNNExplainer"]
class GNNExplainer(nn.Module):
r"""GNNExplainer model from `GNNExplainer: Generating Explanations for
Graph Neural Networks <https://arxiv.org/abs/1903.03894>`__
It identifies compact subgraph structures and small subsets of node features that play a
critical role in GNN-based node classification and graph classification.
To generate an explanation, it learns an edge mask :math:`M` and a feature mask :math:`F`
by optimizing the following objective function.
.. math::
l(y, \hat{y}) + \alpha_1 \|M\|_1 + \alpha_2 H(M) + \beta_1 \|F\|_1 + \beta_2 H(F)
where :math:`l` is the loss function, :math:`y` is the original model prediction,
:math:`\hat{y}` is the model prediction with the edge and feature mask applied, :math:`H` is
the entropy function.
Parameters
----------
model : nn.Module
The GNN model to explain.
* The required arguments of its forward function are graph and feat.
The latter one is for input node features.
* It should also optionally take an eweight argument for edge weights
and multiply the messages by it in message passing.
* The output of its forward function is the logits for the predicted
node/graph classes.
See also the example in :func:`explain_node` and :func:`explain_graph`.
num_hops : int
The number of hops for GNN information aggregation.
lr : float, optional
The learning rate to use, default to 0.01.
num_epochs : int, optional
The number of epochs to train.
alpha1 : float, optional
A higher value will make the explanation edge masks more sparse by decreasing
the sum of the edge mask.
alpha2 : float, optional
A higher value will make the explanation edge masks more sparse by decreasing
the entropy of the edge mask.
beta1 : float, optional
A higher value will make the explanation node feature masks more sparse by
decreasing the mean of the node feature mask.
beta2 : float, optional
A higher value will make the explanation node feature masks more sparse by
decreasing the entropy of the node feature mask.
log : bool, optional
If True, it will log the computation process, default to True.
"""
def __init__(
self,
model,
num_hops,
lr=0.01,
num_epochs=100,
*,
alpha1=0.005,
alpha2=1.0,
beta1=1.0,
beta2=0.1,
log=True,
):
super(GNNExplainer, self).__init__()
self.model = model
self.num_hops = num_hops
self.lr = lr
self.num_epochs = num_epochs
self.alpha1 = alpha1
self.alpha2 = alpha2
self.beta1 = beta1
self.beta2 = beta2
self.log = log
def _init_masks(self, graph, feat):
r"""Initialize learnable feature and edge mask.
Parameters
----------
graph : DGLGraph
Input graph.
feat : Tensor
Input node features.
Returns
-------
feat_mask : Tensor
Feature mask of shape :math:`(1, D)`, where :math:`D`
is the feature size.
edge_mask : Tensor
Edge mask of shape :math:`(E)`, where :math:`E` is the
number of edges.
"""
num_nodes, feat_size = feat.size()
num_edges = graph.num_edges()
device = feat.device
std = 0.1
feat_mask = nn.Parameter(torch.randn(1, feat_size, device=device) * std)
std = nn.init.calculate_gain("relu") * sqrt(2.0 / (2 * num_nodes))
edge_mask = nn.Parameter(torch.randn(num_edges, device=device) * std)
return feat_mask, edge_mask
def _loss_regularize(self, loss, feat_mask, edge_mask):
r"""Add regularization terms to the loss.
Parameters
----------
loss : Tensor
Loss value.
feat_mask : Tensor
Feature mask of shape :math:`(1, D)`, where :math:`D`
is the feature size.
edge_mask : Tensor
Edge mask of shape :math:`(E)`, where :math:`E`
is the number of edges.
Returns
-------
Tensor
Loss value with regularization terms added.
"""
# epsilon for numerical stability
eps = 1e-15
edge_mask = edge_mask.sigmoid()
# Edge mask sparsity regularization
loss = loss + self.alpha1 * torch.sum(edge_mask)
# Edge mask entropy regularization
ent = -edge_mask * torch.log(edge_mask + eps) - (
1 - edge_mask
) * torch.log(1 - edge_mask + eps)
loss = loss + self.alpha2 * ent.mean()
feat_mask = feat_mask.sigmoid()
# Feature mask sparsity regularization
loss = loss + self.beta1 * torch.mean(feat_mask)
# Feature mask entropy regularization
ent = -feat_mask * torch.log(feat_mask + eps) - (
1 - feat_mask
) * torch.log(1 - feat_mask + eps)
loss = loss + self.beta2 * ent.mean()
return loss
def explain_node(self, node_id, graph, feat, **kwargs):
r"""Learn and return a node feature mask and subgraph that play a
crucial role to explain the prediction made by the GNN for node
:attr:`node_id`.
Parameters
----------
node_id : int
The node to explain.
graph : DGLGraph
A homogeneous graph.
feat : Tensor
The input feature of shape :math:`(N, D)`. :math:`N` is the
number of nodes, and :math:`D` is the feature size.
kwargs : dict
Additional arguments passed to the GNN model. Tensors whose
first dimension is the number of nodes or edges will be
assumed to be node/edge features.
Returns
-------
new_node_id : Tensor
The new ID of the input center node.
sg : DGLGraph
The subgraph induced on the k-hop in-neighborhood of the input center node.
feat_mask : Tensor
Learned node feature importance mask of shape :math:`(D)`, where :math:`D` is the
feature size. The values are within range :math:`(0, 1)`.
The higher, the more important.
edge_mask : Tensor
Learned importance mask of the edges in the subgraph, which is a tensor
of shape :math:`(E)`, where :math:`E` is the number of edges in the
subgraph. The values are within range :math:`(0, 1)`.
The higher, the more important.
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch
>>> import torch.nn as nn
>>> from dgl.data import CoraGraphDataset
>>> from dgl.nn import GNNExplainer
>>> # Load dataset
>>> data = CoraGraphDataset()
>>> g = data[0]
>>> features = g.ndata['feat']
>>> labels = g.ndata['label']
>>> train_mask = g.ndata['train_mask']
>>> # Define a model
>>> class Model(nn.Module):
... def __init__(self, in_feats, out_feats):
... super(Model, self).__init__()
... self.linear = nn.Linear(in_feats, out_feats)
...
... def forward(self, graph, feat, eweight=None):
... with graph.local_scope():
... feat = self.linear(feat)
... graph.ndata['h'] = feat
... if eweight is None:
... graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
... else:
... graph.edata['w'] = eweight
... graph.update_all(fn.u_mul_e('h', 'w', 'm'), fn.sum('m', 'h'))
... return graph.ndata['h']
>>> # Train the model
>>> model = Model(features.shape[1], data.num_classes)
>>> criterion = nn.CrossEntropyLoss()
>>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
>>> for epoch in range(10):
... logits = model(g, features)
... loss = criterion(logits[train_mask], labels[train_mask])
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain the prediction for node 10
>>> explainer = GNNExplainer(model, num_hops=1)
>>> new_center, sg, feat_mask, edge_mask = explainer.explain_node(10, g, features)
>>> new_center
tensor([1])
>>> sg.num_edges()
12
>>> # Old IDs of the nodes in the subgraph
>>> sg.ndata[dgl.NID]
tensor([ 9, 10, 11, 12])
>>> # Old IDs of the edges in the subgraph
>>> sg.edata[dgl.EID]
tensor([51, 53, 56, 48, 52, 57, 47, 50, 55, 46, 49, 54])
>>> feat_mask
tensor([0.2638, 0.2738, 0.3039, ..., 0.2794, 0.2643, 0.2733])
>>> edge_mask
tensor([0.0937, 0.1496, 0.8287, 0.8132, 0.8825, 0.8515, 0.8146, 0.0915, 0.1145,
0.9011, 0.1311, 0.8437])
"""
self.model = self.model.to(graph.device)
self.model.eval()
num_nodes = graph.num_nodes()
num_edges = graph.num_edges()
# Extract node-centered k-hop subgraph and
# its associated node and edge features.
sg, inverse_indices = khop_in_subgraph(graph, node_id, self.num_hops)
sg_nodes = sg.ndata[NID].long()
sg_edges = sg.edata[EID].long()
feat = feat[sg_nodes]
for key, item in kwargs.items():
if torch.is_tensor(item) and item.size(0) == num_nodes:
item = item[sg_nodes]
elif torch.is_tensor(item) and item.size(0) == num_edges:
item = item[sg_edges]
kwargs[key] = item
# Get the initial prediction.
with torch.no_grad():
logits = self.model(graph=sg, feat=feat, **kwargs)
pred_label = logits.argmax(dim=-1)
feat_mask, edge_mask = self._init_masks(sg, feat)
params = [feat_mask, edge_mask]
optimizer = torch.optim.Adam(params, lr=self.lr)
if self.log:
pbar = tqdm(total=self.num_epochs)
pbar.set_description(f"Explain node {node_id}")
for _ in range(self.num_epochs):
optimizer.zero_grad()
h = feat * feat_mask.sigmoid()
logits = self.model(
graph=sg, feat=h, eweight=edge_mask.sigmoid(), **kwargs
)
log_probs = logits.log_softmax(dim=-1)
loss = -log_probs[inverse_indices, pred_label[inverse_indices]]
loss = self._loss_regularize(loss, feat_mask, edge_mask)
loss.backward()
optimizer.step()
if self.log:
pbar.update(1)
if self.log:
pbar.close()
feat_mask = feat_mask.detach().sigmoid().squeeze()
edge_mask = edge_mask.detach().sigmoid()
return inverse_indices, sg, feat_mask, edge_mask
def explain_graph(self, graph, feat, **kwargs):
r"""Learn and return a node feature mask and an edge mask that play a
crucial role to explain the prediction made by the GNN for a graph.
Parameters
----------
graph : DGLGraph
A homogeneous graph.
feat : Tensor
The input feature of shape :math:`(N, D)`. :math:`N` is the
number of nodes, and :math:`D` is the feature size.
kwargs : dict
Additional arguments passed to the GNN model. Tensors whose
first dimension is the number of nodes or edges will be
assumed to be node/edge features.
Returns
-------
feat_mask : Tensor
Learned feature importance mask of shape :math:`(D)`, where :math:`D` is the
feature size. The values are within range :math:`(0, 1)`.
The higher, the more important.
edge_mask : Tensor
Learned importance mask of the edges in the graph, which is a tensor
of shape :math:`(E)`, where :math:`E` is the number of edges in the
graph. The values are within range :math:`(0, 1)`. The higher,
the more important.
Examples
--------
>>> import dgl.function as fn
>>> import torch
>>> import torch.nn as nn
>>> from dgl.data import GINDataset
>>> from dgl.dataloading import GraphDataLoader
>>> from dgl.nn import AvgPooling, GNNExplainer
>>> # Load dataset
>>> data = GINDataset('MUTAG', self_loop=True)
>>> dataloader = GraphDataLoader(data, batch_size=64, shuffle=True)
>>> # Define a model
>>> class Model(nn.Module):
... def __init__(self, in_feats, out_feats):
... super(Model, self).__init__()
... self.linear = nn.Linear(in_feats, out_feats)
... self.pool = AvgPooling()
...
... def forward(self, graph, feat, eweight=None):
... with graph.local_scope():
... feat = self.linear(feat)
... graph.ndata['h'] = feat
... if eweight is None:
... graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
... else:
... graph.edata['w'] = eweight
... graph.update_all(fn.u_mul_e('h', 'w', 'm'), fn.sum('m', 'h'))
... return self.pool(graph, graph.ndata['h'])
>>> # Train the model
>>> feat_size = data[0][0].ndata['attr'].shape[1]
>>> model = Model(feat_size, data.gclasses)
>>> criterion = nn.CrossEntropyLoss()
>>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
>>> for bg, labels in dataloader:
... logits = model(bg, bg.ndata['attr'])
... loss = criterion(logits, labels)
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain the prediction for graph 0
>>> explainer = GNNExplainer(model, num_hops=1)
>>> g, _ = data[0]
>>> features = g.ndata['attr']
>>> feat_mask, edge_mask = explainer.explain_graph(g, features)
>>> feat_mask
tensor([0.2362, 0.2497, 0.2622, 0.2675, 0.2649, 0.2962, 0.2533])
>>> edge_mask
tensor([0.2154, 0.2235, 0.8325, ..., 0.7787, 0.1735, 0.1847])
"""
self.model = self.model.to(graph.device)
self.model.eval()
# Get the initial prediction.
with torch.no_grad():
logits = self.model(graph=graph, feat=feat, **kwargs)
pred_label = logits.argmax(dim=-1)
feat_mask, edge_mask = self._init_masks(graph, feat)
params = [feat_mask, edge_mask]
optimizer = torch.optim.Adam(params, lr=self.lr)
if self.log:
pbar = tqdm(total=self.num_epochs)
pbar.set_description("Explain graph")
for _ in range(self.num_epochs):
optimizer.zero_grad()
h = feat * feat_mask.sigmoid()
logits = self.model(
graph=graph, feat=h, eweight=edge_mask.sigmoid(), **kwargs
)
log_probs = logits.log_softmax(dim=-1)
loss = -log_probs[0, pred_label[0]]
loss = self._loss_regularize(loss, feat_mask, edge_mask)
loss.backward()
optimizer.step()
if self.log:
pbar.update(1)
if self.log:
pbar.close()
feat_mask = feat_mask.detach().sigmoid().squeeze()
edge_mask = edge_mask.detach().sigmoid()
return feat_mask, edge_mask
class HeteroGNNExplainer(nn.Module):
r"""GNNExplainer model from `GNNExplainer: Generating Explanations for
Graph Neural Networks <https://arxiv.org/abs/1903.03894>`__, adapted for heterogeneous graphs
It identifies compact subgraph structures and small subsets of node features that play a
critical role in GNN-based node classification and graph classification.
To generate an explanation, it learns an edge mask :math:`M` and a feature mask :math:`F`
by optimizing the following objective function.
.. math::
l(y, \hat{y}) + \alpha_1 \|M\|_1 + \alpha_2 H(M) + \beta_1 \|F\|_1 + \beta_2 H(F)
where :math:`l` is the loss function, :math:`y` is the original model prediction,
:math:`\hat{y}` is the model prediction with the edge and feature mask applied, :math:`H` is
the entropy function.
Parameters
----------
model : nn.Module
The GNN model to explain.
* The required arguments of its forward function are graph and feat.
The latter one is for input node features.
* It should also optionally take an eweight argument for edge weights
and multiply the messages by it in message passing.
* The output of its forward function is the logits for the predicted
node/graph classes.
See also the example in :func:`explain_node` and :func:`explain_graph`.
num_hops : int
The number of hops for GNN information aggregation.
lr : float, optional
The learning rate to use, default to 0.01.
num_epochs : int, optional
The number of epochs to train.
alpha1 : float, optional
A higher value will make the explanation edge masks more sparse by decreasing
the sum of the edge mask.
alpha2 : float, optional
A higher value will make the explanation edge masks more sparse by decreasing
the entropy of the edge mask.
beta1 : float, optional
A higher value will make the explanation node feature masks more sparse by
decreasing the mean of the node feature mask.
beta2 : float, optional
A higher value will make the explanation node feature masks more sparse by
decreasing the entropy of the node feature mask.
log : bool, optional
If True, it will log the computation process, default to True.
"""
def __init__(
self,
model,
num_hops,
lr=0.01,
num_epochs=100,
*,
alpha1=0.005,
alpha2=1.0,
beta1=1.0,
beta2=0.1,
log=True,
):
super(HeteroGNNExplainer, self).__init__()
self.model = model
self.num_hops = num_hops
self.lr = lr
self.num_epochs = num_epochs
self.alpha1 = alpha1
self.alpha2 = alpha2
self.beta1 = beta1
self.beta2 = beta2
self.log = log
def _init_masks(self, graph, feat):
r"""Initialize learnable feature and edge mask.
Parameters
----------
graph : DGLGraph
Input graph.
feat : dict[str, Tensor]
The dictionary that associates input node features (values) with
the respective node types (keys) present in the graph.
Returns
-------
feat_masks : dict[str, Tensor]
The dictionary that associates the node feature masks (values) with
the respective node types (keys). The feature masks are of shape :math:`(1, D_t)`,
where :math:`D_t` is the feature size for node type :math:`t`.
edge_masks : dict[tuple[str], Tensor]
The dictionary that associates the edge masks (values) with
the respective canonical edge types (keys). The edge masks are of shape :math:`(E_t)`,
where :math:`E_t` is the number of edges for canonical edge type :math:`t`.
"""
device = graph.device
feat_masks = {}
std = 0.1
for node_type, feature in feat.items():
_, feat_size = feature.size()
feat_masks[node_type] = nn.Parameter(
torch.randn(1, feat_size, device=device) * std
)
edge_masks = {}
for canonical_etype in graph.canonical_etypes:
src_num_nodes = graph.num_nodes(canonical_etype[0])
dst_num_nodes = graph.num_nodes(canonical_etype[-1])
num_nodes_sum = src_num_nodes + dst_num_nodes
num_edges = graph.num_edges(canonical_etype)
std = nn.init.calculate_gain("relu")
if num_nodes_sum > 0:
std *= sqrt(2.0 / num_nodes_sum)
edge_masks[canonical_etype] = nn.Parameter(
torch.randn(num_edges, device=device) * std
)
return feat_masks, edge_masks
def _loss_regularize(self, loss, feat_masks, edge_masks):
r"""Add regularization terms to the loss.
Parameters
----------
loss : Tensor
Loss value.
feat_masks : dict[str, Tensor]
The dictionary that associates the node feature masks (values) with
the respective node types (keys). The feature masks are of shape :math:`(1, D_t)`,
where :math:`D_t` is the feature size for node type :math:`t`.
edge_masks : dict[tuple[str], Tensor]
The dictionary that associates the edge masks (values) with
the respective canonical edge types (keys). The edge masks are of shape :math:`(E_t)`,
where :math:`E_t` is the number of edges for canonical edge type :math:`t`.
Returns
-------
Tensor
Loss value with regularization terms added.
"""
# epsilon for numerical stability
eps = 1e-15
for edge_mask in edge_masks.values():
edge_mask = edge_mask.sigmoid()
# Edge mask sparsity regularization
loss = loss + self.alpha1 * torch.sum(edge_mask)
# Edge mask entropy regularization
ent = -edge_mask * torch.log(edge_mask + eps) - (
1 - edge_mask
) * torch.log(1 - edge_mask + eps)
loss = loss + self.alpha2 * ent.mean()
for feat_mask in feat_masks.values():
feat_mask = feat_mask.sigmoid()
# Feature mask sparsity regularization
loss = loss + self.beta1 * torch.mean(feat_mask)
# Feature mask entropy regularization
ent = -feat_mask * torch.log(feat_mask + eps) - (
1 - feat_mask
) * torch.log(1 - feat_mask + eps)
loss = loss + self.beta2 * ent.mean()
return loss
def explain_node(self, ntype, node_id, graph, feat, **kwargs):
r"""Learn and return node feature masks and a subgraph that play a
crucial role to explain the prediction made by the GNN for node
:attr:`node_id` of type :attr:`ntype`.
It requires :attr:`model` to return a dictionary mapping node types to type-specific
predictions.
Parameters
----------
ntype : str
The type of the node to explain. :attr:`model` must be trained to
make predictions for this particular node type.
node_id : int
The ID of the node to explain.
graph : DGLGraph
A heterogeneous graph.
feat : dict[str, Tensor]
The dictionary that associates input node features (values) with
the respective node types (keys) present in the graph.
The input features are of shape :math:`(N_t, D_t)`. :math:`N_t` is the
number of nodes for node type :math:`t`, and :math:`D_t` is the feature size for
node type :math:`t`
kwargs : dict
Additional arguments passed to the GNN model.
Returns
-------
new_node_id : Tensor
The new ID of the input center node.
sg : DGLGraph
The subgraph induced on the k-hop in-neighborhood of the input center node.
feat_mask : dict[str, Tensor]
The dictionary that associates the learned node feature importance masks (values) with
the respective node types (keys). The masks are of shape :math:`(D_t)`, where
:math:`D_t` is the node feature size for node type :attr:`t`. The values are within
range :math:`(0, 1)`. The higher, the more important.
edge_mask : dict[Tuple[str], Tensor]
The dictionary that associates the learned edge importance masks (values) with
the respective canonical edge types (keys). The masks are of shape :math:`(E_t)`,
where :math:`E_t` is the number of edges for canonical edge type :math:`t` in the
subgraph. The values are within range :math:`(0, 1)`.
The higher, the more important.
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch as th
>>> import torch.nn as nn
>>> import torch.nn.functional as F
>>> from dgl.nn import HeteroGNNExplainer
>>> class Model(nn.Module):
... def __init__(self, in_dim, num_classes, canonical_etypes):
... super(Model, self).__init__()
... self.etype_weights = nn.ModuleDict({
... '_'.join(c_etype): nn.Linear(in_dim, num_classes)
... for c_etype in canonical_etypes
... })
...
... def forward(self, graph, feat, eweight=None):
... with graph.local_scope():
... c_etype_func_dict = {}
... for c_etype in graph.canonical_etypes:
... src_type, etype, dst_type = c_etype
... wh = self.etype_weights['_'.join(c_etype)](feat[src_type])
... graph.nodes[src_type].data[f'h_{c_etype}'] = wh
... if eweight is None:
... c_etype_func_dict[c_etype] = (fn.copy_u(f'h_{c_etype}', 'm'),
... fn.mean('m', 'h'))
... else:
... graph.edges[c_etype].data['w'] = eweight[c_etype]
... c_etype_func_dict[c_etype] = (
... fn.u_mul_e(f'h_{c_etype}', 'w', 'm'), fn.mean('m', 'h'))
... graph.multi_update_all(c_etype_func_dict, 'sum')
... return graph.ndata['h']
>>> input_dim = 5
>>> num_classes = 2
>>> g = dgl.heterograph({
... ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1])})
>>> g.nodes['user'].data['h'] = th.randn(g.num_nodes('user'), input_dim)
>>> g.nodes['game'].data['h'] = th.randn(g.num_nodes('game'), input_dim)
>>> transform = dgl.transforms.AddReverse()
>>> g = transform(g)
>>> # define and train the model
>>> model = Model(input_dim, num_classes, g.canonical_etypes)
>>> feat = g.ndata['h']
>>> optimizer = th.optim.Adam(model.parameters())
>>> for epoch in range(10):
... logits = model(g, feat)['user']
... loss = F.cross_entropy(logits, th.tensor([1, 1, 1]))
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain the prediction for node 0 of type 'user'
>>> explainer = HeteroGNNExplainer(model, num_hops=1)
>>> new_center, sg, feat_mask, edge_mask = explainer.explain_node('user', 0, g, feat)
>>> new_center
tensor([0])
>>> sg
Graph(num_nodes={'game': 1, 'user': 1},
num_edges={('game', 'rev_plays', 'user'): 1, ('user', 'plays', 'game'): 1,
('user', 'rev_rev_plays', 'game'): 1},
metagraph=[('game', 'user', 'rev_plays'), ('user', 'game', 'plays'),
('user', 'game', 'rev_rev_plays')])
>>> feat_mask
{'game': tensor([0.2348, 0.2780, 0.2611, 0.2513, 0.2823]),
'user': tensor([0.2716, 0.2450, 0.2658, 0.2876, 0.2738])}
>>> edge_mask
{('game', 'rev_plays', 'user'): tensor([0.0630]),
('user', 'plays', 'game'): tensor([0.1939]),
('user', 'rev_rev_plays', 'game'): tensor([0.9166])}
"""
self.model = self.model.to(graph.device)
self.model.eval()
# Extract node-centered k-hop subgraph and
# its associated node and edge features.
sg, inverse_indices = khop_in_subgraph(
graph, {ntype: node_id}, self.num_hops
)
inverse_indices = inverse_indices[ntype]
sg_nodes = sg.ndata[NID]
sg_feat = {}
for node_type in sg_nodes.keys():
sg_feat[node_type] = feat[node_type][sg_nodes[node_type].long()]
# Get the initial prediction.
with torch.no_grad():
logits = self.model(graph=sg, feat=sg_feat, **kwargs)[ntype]
pred_label = logits.argmax(dim=-1)
feat_mask, edge_mask = self._init_masks(sg, sg_feat)
params = [*feat_mask.values(), *edge_mask.values()]
optimizer = torch.optim.Adam(params, lr=self.lr)
if self.log:
pbar = tqdm(total=self.num_epochs)
pbar.set_description(f"Explain node {node_id} with type {ntype}")
for _ in range(self.num_epochs):
optimizer.zero_grad()
h = {}
for node_type, sg_node_feat in sg_feat.items():
h[node_type] = sg_node_feat * feat_mask[node_type].sigmoid()
eweight = {}
for canonical_etype, canonical_etype_mask in edge_mask.items():
eweight[canonical_etype] = canonical_etype_mask.sigmoid()
logits = self.model(graph=sg, feat=h, eweight=eweight, **kwargs)[
ntype
]
log_probs = logits.log_softmax(dim=-1)
loss = -log_probs[inverse_indices, pred_label[inverse_indices]]
loss = self._loss_regularize(loss, feat_mask, edge_mask)
loss.backward()
optimizer.step()
if self.log:
pbar.update(1)
if self.log:
pbar.close()
for node_type in feat_mask:
feat_mask[node_type] = (
feat_mask[node_type].detach().sigmoid().squeeze()
)
for canonical_etype in edge_mask:
edge_mask[canonical_etype] = (
edge_mask[canonical_etype].detach().sigmoid()
)
return inverse_indices, sg, feat_mask, edge_mask
def explain_graph(self, graph, feat, **kwargs):
r"""Learn and return node feature masks and edge masks that play a
crucial role to explain the prediction made by the GNN for a graph.
Parameters
----------
graph : DGLGraph
A heterogeneous graph that will be explained.
feat : dict[str, Tensor]
The dictionary that associates input node features (values) with
the respective node types (keys) present in the graph.
The input features are of shape :math:`(N_t, D_t)`. :math:`N_t` is the
number of nodes for node type :math:`t`, and :math:`D_t` is the feature size for
node type :math:`t`
kwargs : dict
Additional arguments passed to the GNN model.
Returns
-------
feat_mask : dict[str, Tensor]
The dictionary that associates the learned node feature importance masks (values) with
the respective node types (keys). The masks are of shape :math:`(D_t)`, where
:math:`D_t` is the node feature size for node type :attr:`t`. The values are within
range :math:`(0, 1)`. The higher, the more important.
edge_mask : dict[Tuple[str], Tensor]
The dictionary that associates the learned edge importance masks (values) with
the respective canonical edge types (keys). The masks are of shape :math:`(E_t)`,
where :math:`E_t` is the number of edges for canonical edge type :math:`t` in the
graph. The values are within range :math:`(0, 1)`. The higher, the more important.
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch as th
>>> import torch.nn as nn
>>> import torch.nn.functional as F
>>> from dgl.nn import HeteroGNNExplainer
>>> class Model(nn.Module):
... def __init__(self, in_dim, num_classes, canonical_etypes):
... super(Model, self).__init__()
... self.etype_weights = nn.ModuleDict({
... '_'.join(c_etype): nn.Linear(in_dim, num_classes)
... for c_etype in canonical_etypes
... })
...
... def forward(self, graph, feat, eweight=None):
... with graph.local_scope():
... c_etype_func_dict = {}
... for c_etype in graph.canonical_etypes:
... src_type, etype, dst_type = c_etype
... wh = self.etype_weights['_'.join(c_etype)](feat[src_type])
... graph.nodes[src_type].data[f'h_{c_etype}'] = wh
... if eweight is None:
... c_etype_func_dict[c_etype] = (fn.copy_u(f'h_{c_etype}', 'm'),
... fn.mean('m', 'h'))
... else:
... graph.edges[c_etype].data['w'] = eweight[c_etype]
... c_etype_func_dict[c_etype] = (
... fn.u_mul_e(f'h_{c_etype}', 'w', 'm'), fn.mean('m', 'h'))
... graph.multi_update_all(c_etype_func_dict, 'sum')
... hg = 0
... for ntype in graph.ntypes:
... if graph.num_nodes(ntype):
... hg = hg + dgl.mean_nodes(graph, 'h', ntype=ntype)
... return hg
>>> input_dim = 5
>>> num_classes = 2
>>> g = dgl.heterograph({
... ('user', 'plays', 'game'): ([0, 1, 1, 2], [0, 0, 1, 1])})
>>> g.nodes['user'].data['h'] = th.randn(g.num_nodes('user'), input_dim)
>>> g.nodes['game'].data['h'] = th.randn(g.num_nodes('game'), input_dim)
>>> transform = dgl.transforms.AddReverse()
>>> g = transform(g)
>>> # define and train the model
>>> model = Model(input_dim, num_classes, g.canonical_etypes)
>>> feat = g.ndata['h']
>>> optimizer = th.optim.Adam(model.parameters())
>>> for epoch in range(10):
... logits = model(g, feat)
... loss = F.cross_entropy(logits, th.tensor([1]))
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain for the graph
>>> explainer = HeteroGNNExplainer(model, num_hops=1)
>>> feat_mask, edge_mask = explainer.explain_graph(g, feat)
>>> feat_mask
{'game': tensor([0.2684, 0.2597, 0.3135, 0.2976, 0.2607]),
'user': tensor([0.2216, 0.2908, 0.2644, 0.2738, 0.2663])}
>>> edge_mask
{('game', 'rev_plays', 'user'): tensor([0.8922, 0.1966, 0.8371, 0.1330]),
('user', 'plays', 'game'): tensor([0.1785, 0.1696, 0.8065, 0.2167])}
"""
self.model = self.model.to(graph.device)
self.model.eval()
# Get the initial prediction.
with torch.no_grad():
logits = self.model(graph=graph, feat=feat, **kwargs)
pred_label = logits.argmax(dim=-1)
feat_mask, edge_mask = self._init_masks(graph, feat)
params = [*feat_mask.values(), *edge_mask.values()]
optimizer = torch.optim.Adam(params, lr=self.lr)
if self.log:
pbar = tqdm(total=self.num_epochs)
pbar.set_description("Explain graph")
for _ in range(self.num_epochs):
optimizer.zero_grad()
h = {}
for node_type, node_feat in feat.items():
h[node_type] = node_feat * feat_mask[node_type].sigmoid()
eweight = {}
for canonical_etype, canonical_etype_mask in edge_mask.items():
eweight[canonical_etype] = canonical_etype_mask.sigmoid()
logits = self.model(graph=graph, feat=h, eweight=eweight, **kwargs)
log_probs = logits.log_softmax(dim=-1)
loss = -log_probs[0, pred_label[0]]
loss = self._loss_regularize(loss, feat_mask, edge_mask)
loss.backward()
optimizer.step()
if self.log:
pbar.update(1)
if self.log:
pbar.close()
for node_type in feat_mask:
feat_mask[node_type] = (
feat_mask[node_type].detach().sigmoid().squeeze()
)
for canonical_etype in edge_mask:
edge_mask[canonical_etype] = (
edge_mask[canonical_etype].detach().sigmoid()
)
return feat_mask, edge_mask
python/dgl/nn/pytorch/explain/pgexplainer.py
+1139
View File
File diff suppressed because it is too large Load Diff
python/dgl/nn/pytorch/explain/subgraphx.py
+807
View File
@@ -0,0 +1,807 @@
"""Torch Module for SubgraphX"""
import math
import networkx as nx
import numpy as np
import torch
import torch.nn as nn
from .... import to_heterogeneous, to_homogeneous
from ....base import NID
from ....convert import to_networkx
from ....subgraph import node_subgraph
from ....transforms.functional import remove_nodes
__all__ = ["SubgraphX", "HeteroSubgraphX"]
class MCTSNode:
r"""Monte Carlo Tree Search Node
Parameters
----------
nodes : Tensor
The node IDs of the graph that are associated with this tree node
"""
def __init__(self, nodes):
self.nodes = nodes
self.num_visit = 0
self.total_reward = 0.0
self.immediate_reward = 0.0
self.children = []
def __repr__(self):
r"""Get the string representation of the node.
Returns
-------
str
The string representation of the node
"""
return str(self.nodes)
class SubgraphX(nn.Module):
r"""SubgraphX from `On Explainability of Graph Neural Networks via Subgraph
Explorations <https://arxiv.org/abs/2102.05152>`
It identifies the most important subgraph from the original graph that
plays a critical role in GNN-based graph classification.
It employs Monte Carlo tree search (MCTS) in efficiently exploring
different subgraphs for explanation and uses Shapley values as the measure
of subgraph importance.
Parameters
----------
model : nn.Module
The GNN model to explain that tackles multiclass graph classification
* Its forward function must have the form
:attr:`forward(self, graph, nfeat)`.
* The output of its forward function is the logits.
num_hops : int
Number of message passing layers in the model
coef : float, optional
This hyperparameter controls the trade-off between exploration and
exploitation. A higher value encourages the algorithm to explore
relatively unvisited nodes. Default: 10.0
high2low : bool, optional
If True, it will use the "High2low" strategy for pruning actions,
expanding children nodes from high degree to low degree when extending
the children nodes in the search tree. Otherwise, it will use the
"Low2high" strategy. Default: True
num_child : int, optional
This is the number of children nodes to expand when extending the
children nodes in the search tree. Default: 12
num_rollouts : int, optional
This is the number of rollouts for MCTS. Default: 20
node_min : int, optional
This is the threshold to define a leaf node based on the number of
nodes in a subgraph. Default: 3
shapley_steps : int, optional
This is the number of steps for Monte Carlo sampling in estimating
Shapley values. Default: 100
log : bool, optional
If True, it will log the progress. Default: False
"""
def __init__(
self,
model,
num_hops,
coef=10.0,
high2low=True,
num_child=12,
num_rollouts=20,
node_min=3,
shapley_steps=100,
log=False,
):
super().__init__()
self.num_hops = num_hops
self.coef = coef
self.high2low = high2low
self.num_child = num_child
self.num_rollouts = num_rollouts
self.node_min = node_min
self.shapley_steps = shapley_steps
self.log = log
self.model = model
def shapley(self, subgraph_nodes):
r"""Compute Shapley value with Monte Carlo approximation.
Parameters
----------
subgraph_nodes : tensor
The tensor node ids of the subgraph that are associated with this
tree node
Returns
-------
float
Shapley value
"""
num_nodes = self.graph.num_nodes()
subgraph_nodes = subgraph_nodes.tolist()
# Obtain neighboring nodes of the subgraph g_i, P'.
local_region = subgraph_nodes
for _ in range(self.num_hops - 1):
in_neighbors, _ = self.graph.in_edges(local_region)
_, out_neighbors = self.graph.out_edges(local_region)
neighbors = torch.cat([in_neighbors, out_neighbors]).tolist()
local_region = list(set(local_region + neighbors))
split_point = num_nodes
coalition_space = list(set(local_region) - set(subgraph_nodes)) + [
split_point
]
marginal_contributions = []
device = self.feat.device
for _ in range(self.shapley_steps):
permuted_space = np.random.permutation(coalition_space)
split_idx = int(np.where(permuted_space == split_point)[0])
selected_nodes = permuted_space[:split_idx]
# Mask for coalition set S_i
exclude_mask = torch.ones(num_nodes)
exclude_mask[local_region] = 0.0
exclude_mask[selected_nodes] = 1.0
# Mask for set S_i and g_i
include_mask = exclude_mask.clone()
include_mask[subgraph_nodes] = 1.0
exclude_feat = self.feat * exclude_mask.unsqueeze(1).to(device)
include_feat = self.feat * include_mask.unsqueeze(1).to(device)
with torch.no_grad():
exclude_probs = self.model(
self.graph, exclude_feat, **self.kwargs
).softmax(dim=-1)
exclude_value = exclude_probs[:, self.target_class]
include_probs = self.model(
self.graph, include_feat, **self.kwargs
).softmax(dim=-1)
include_value = include_probs[:, self.target_class]
marginal_contributions.append(include_value - exclude_value)
return torch.cat(marginal_contributions).mean().item()
def get_mcts_children(self, mcts_node):
r"""Get the children of the MCTS node for the search.
Parameters
----------
mcts_node : MCTSNode
Node in MCTS
Returns
-------
list
Children nodes after pruning
"""
if len(mcts_node.children) > 0:
return mcts_node.children
subg = node_subgraph(self.graph, mcts_node.nodes)
node_degrees = subg.out_degrees() + subg.in_degrees()
k = min(subg.num_nodes(), self.num_child)
chosen_nodes = torch.topk(
node_degrees, k, largest=self.high2low
).indices
mcts_children_maps = dict()
for node in chosen_nodes:
new_subg = remove_nodes(subg, node.to(subg.idtype), store_ids=True)
# Get the largest weakly connected component in the subgraph.
nx_graph = to_networkx(new_subg.cpu())
largest_cc_nids = list(
max(nx.weakly_connected_components(nx_graph), key=len)
)
# Map to the original node IDs.
largest_cc_nids = new_subg.ndata[NID][largest_cc_nids].long()
largest_cc_nids = subg.ndata[NID][largest_cc_nids].sort().values
if str(largest_cc_nids) not in self.mcts_node_maps:
child_mcts_node = MCTSNode(largest_cc_nids)
self.mcts_node_maps[str(child_mcts_node)] = child_mcts_node
else:
child_mcts_node = self.mcts_node_maps[str(largest_cc_nids)]
if str(child_mcts_node) not in mcts_children_maps:
mcts_children_maps[str(child_mcts_node)] = child_mcts_node
mcts_node.children = list(mcts_children_maps.values())
for child_mcts_node in mcts_node.children:
if child_mcts_node.immediate_reward == 0:
child_mcts_node.immediate_reward = self.shapley(
child_mcts_node.nodes
)
return mcts_node.children
def mcts_rollout(self, mcts_node):
r"""Perform a MCTS rollout.
Parameters
----------
mcts_node : MCTSNode
Starting node for MCTS
Returns
-------
float
Reward for visiting the node this time
"""
if len(mcts_node.nodes) <= self.node_min:
return mcts_node.immediate_reward
children_nodes = self.get_mcts_children(mcts_node)
children_visit_sum = sum([child.num_visit for child in children_nodes])
children_visit_sum_sqrt = math.sqrt(children_visit_sum)
chosen_child = max(
children_nodes,
key=lambda c: c.total_reward / max(c.num_visit, 1)
+ self.coef
* c.immediate_reward
* children_visit_sum_sqrt
/ (1 + c.num_visit),
)
reward = self.mcts_rollout(chosen_child)
chosen_child.num_visit += 1
chosen_child.total_reward += reward
return reward
def explain_graph(self, graph, feat, target_class, **kwargs):
r"""Find the most important subgraph from the original graph for the
model to classify the graph into the target class.
Parameters
----------
graph : DGLGraph
A homogeneous graph
feat : Tensor
The input node feature of shape :math:`(N, D)`, :math:`N` is the
number of nodes, and :math:`D` is the feature size
target_class : int
The target class to explain
kwargs : dict
Additional arguments passed to the GNN model
Returns
-------
Tensor
Nodes that represent the most important subgraph
Examples
--------
>>> import torch
>>> import torch.nn as nn
>>> import torch.nn.functional as F
>>> from dgl.data import GINDataset
>>> from dgl.dataloading import GraphDataLoader
>>> from dgl.nn import GraphConv, AvgPooling, SubgraphX
>>> # Define the model
>>> class Model(nn.Module):
... def __init__(self, in_dim, n_classes, hidden_dim=128):
... super().__init__()
... self.conv1 = GraphConv(in_dim, hidden_dim)
... self.conv2 = GraphConv(hidden_dim, n_classes)
... self.pool = AvgPooling()
...
... def forward(self, g, h):
... h = F.relu(self.conv1(g, h))
... h = self.conv2(g, h)
... return self.pool(g, h)
>>> # Load dataset
>>> data = GINDataset('MUTAG', self_loop=True)
>>> dataloader = GraphDataLoader(data, batch_size=64, shuffle=True)
>>> # Train the model
>>> feat_size = data[0][0].ndata['attr'].shape[1]
>>> model = Model(feat_size, data.gclasses)
>>> criterion = nn.CrossEntropyLoss()
>>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
>>> for bg, labels in dataloader:
... logits = model(bg, bg.ndata['attr'])
... loss = criterion(logits, labels)
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Initialize the explainer
>>> explainer = SubgraphX(model, num_hops=2)
>>> # Explain the prediction for graph 0
>>> graph, l = data[0]
>>> graph_feat = graph.ndata.pop("attr")
>>> g_nodes_explain = explainer.explain_graph(graph, graph_feat,
... target_class=l)
"""
self.model.eval()
assert (
graph.num_nodes() > self.node_min
), f"The number of nodes in the\
graph {graph.num_nodes()} should be bigger than {self.node_min}."
self.graph = graph
self.feat = feat
self.target_class = target_class
self.kwargs = kwargs
# book all nodes in MCTS
self.mcts_node_maps = dict()
root = MCTSNode(graph.nodes())
self.mcts_node_maps[str(root)] = root
for i in range(self.num_rollouts):
if self.log:
print(
f"Rollout {i}/{self.num_rollouts}, \
{len(self.mcts_node_maps)} subgraphs have been explored."
)
self.mcts_rollout(root)
best_leaf = None
best_immediate_reward = float("-inf")
for mcts_node in self.mcts_node_maps.values():
if len(mcts_node.nodes) > self.node_min:
continue
if mcts_node.immediate_reward > best_immediate_reward:
best_leaf = mcts_node
best_immediate_reward = best_leaf.immediate_reward
return best_leaf.nodes
class HeteroSubgraphX(nn.Module):
r"""SubgraphX from `On Explainability of Graph Neural Networks via Subgraph
Explorations <https://arxiv.org/abs/2102.05152>`__, adapted for heterogeneous graphs
It identifies the most important subgraph from the original graph that
plays a critical role in GNN-based graph classification.
It employs Monte Carlo tree search (MCTS) in efficiently exploring
different subgraphs for explanation and uses Shapley values as the measure
of subgraph importance.
Parameters
----------
model : nn.Module
The GNN model to explain that tackles multiclass graph classification
* Its forward function must have the form
:attr:`forward(self, graph, nfeat)`.
* The output of its forward function is the logits.
num_hops : int
Number of message passing layers in the model
coef : float, optional
This hyperparameter controls the trade-off between exploration and
exploitation. A higher value encourages the algorithm to explore
relatively unvisited nodes. Default: 10.0
high2low : bool, optional
If True, it will use the "High2low" strategy for pruning actions,
expanding children nodes from high degree to low degree when extending
the children nodes in the search tree. Otherwise, it will use the
"Low2high" strategy. Default: True
num_child : int, optional
This is the number of children nodes to expand when extending the
children nodes in the search tree. Default: 12
num_rollouts : int, optional
This is the number of rollouts for MCTS. Default: 20
node_min : int, optional
This is the threshold to define a leaf node based on the number of
nodes in a subgraph. Default: 3
shapley_steps : int, optional
This is the number of steps for Monte Carlo sampling in estimating
Shapley values. Default: 100
log : bool, optional
If True, it will log the progress. Default: False
"""
def __init__(
self,
model,
num_hops,
coef=10.0,
high2low=True,
num_child=12,
num_rollouts=20,
node_min=3,
shapley_steps=100,
log=False,
):
super().__init__()
self.num_hops = num_hops
self.coef = coef
self.high2low = high2low
self.num_child = num_child
self.num_rollouts = num_rollouts
self.node_min = node_min
self.shapley_steps = shapley_steps
self.log = log
self.model = model
def shapley(self, subgraph_nodes):
r"""Compute Shapley value with Monte Carlo approximation.
Parameters
----------
subgraph_nodes : dict[str, Tensor]
subgraph_nodes[nty] gives the tensor node IDs of node type nty
in the subgraph, which are associated with this tree node
Returns
-------
float
Shapley value
"""
# Obtain neighboring nodes of the subgraph g_i, P'.
local_regions = {
ntype: nodes.tolist() for ntype, nodes in subgraph_nodes.items()
}
for _ in range(self.num_hops - 1):
for c_etype in self.graph.canonical_etypes:
src_ntype, _, dst_ntype = c_etype
if (
src_ntype not in local_regions
or dst_ntype not in local_regions
):
continue
in_neighbors, _ = self.graph.in_edges(
local_regions[dst_ntype], etype=c_etype
)
_, out_neighbors = self.graph.out_edges(
local_regions[src_ntype], etype=c_etype
)
local_regions[src_ntype] = list(
set(local_regions[src_ntype] + in_neighbors.tolist())
)
local_regions[dst_ntype] = list(
set(local_regions[dst_ntype] + out_neighbors.tolist())
)
split_point = self.graph.num_nodes()
coalition_space = {
ntype: list(
set(local_regions[ntype]) - set(subgraph_nodes[ntype].tolist())
)
+ [split_point]
for ntype in subgraph_nodes.keys()
}
marginal_contributions = []
for _ in range(self.shapley_steps):
selected_node_map = dict()
for ntype, nodes in coalition_space.items():
permuted_space = np.random.permutation(nodes)
split_idx = int(np.where(permuted_space == split_point)[0])
selected_node_map[ntype] = permuted_space[:split_idx]
# Mask for coalition set S_i
exclude_mask = {
ntype: torch.ones(self.graph.num_nodes(ntype))
for ntype in self.graph.ntypes
}
for ntype, region in local_regions.items():
exclude_mask[ntype][region] = 0.0
for ntype, selected_nodes in selected_node_map.items():
exclude_mask[ntype][selected_nodes] = 1.0
# Mask for set S_i and g_i
include_mask = {
ntype: exclude_mask[ntype].clone()
for ntype in self.graph.ntypes
}
for ntype, subgn in subgraph_nodes.items():
exclude_mask[ntype][subgn] = 1.0
exclude_feat = {
ntype: self.feat[ntype]
* exclude_mask[ntype].unsqueeze(1).to(self.feat[ntype].device)
for ntype in self.graph.ntypes
}
include_feat = {
ntype: self.feat[ntype]
* include_mask[ntype].unsqueeze(1).to(self.feat[ntype].device)
for ntype in self.graph.ntypes
}
with torch.no_grad():
exclude_probs = self.model(
self.graph, exclude_feat, **self.kwargs
).softmax(dim=-1)
exclude_value = exclude_probs[:, self.target_class]
include_probs = self.model(
self.graph, include_feat, **self.kwargs
).softmax(dim=-1)
include_value = include_probs[:, self.target_class]
marginal_contributions.append(include_value - exclude_value)
return torch.cat(marginal_contributions).mean().item()
def get_mcts_children(self, mcts_node):
r"""Get the children of the MCTS node for the search.
Parameters
----------
mcts_node : MCTSNode
Node in MCTS
Returns
-------
list
Children nodes after pruning
"""
if len(mcts_node.children) > 0:
return mcts_node.children
subg = node_subgraph(self.graph, mcts_node.nodes)
# Choose k nodes based on the highest degree in the subgraph
node_degrees_map = {
ntype: torch.zeros(
subg.num_nodes(ntype), device=subg.nodes(ntype).device
)
for ntype in subg.ntypes
}
for c_etype in subg.canonical_etypes:
src_ntype, _, dst_ntype = c_etype
node_degrees_map[src_ntype] += subg.out_degrees(etype=c_etype)
node_degrees_map[dst_ntype] += subg.in_degrees(etype=c_etype)
node_degrees_list = [
((ntype, i), degree)
for ntype, node_degrees in node_degrees_map.items()
for i, degree in enumerate(node_degrees)
]
node_degrees = torch.stack([v for _, v in node_degrees_list])
k = min(subg.num_nodes(), self.num_child)
chosen_node_indicies = torch.topk(
node_degrees, k, largest=self.high2low
).indices
chosen_nodes = [node_degrees_list[i][0] for i in chosen_node_indicies]
mcts_children_maps = dict()
for ntype, node in chosen_nodes:
new_subg = remove_nodes(subg, node, ntype, store_ids=True)
if new_subg.num_edges() > 0:
new_subg_homo = to_homogeneous(new_subg)
# Get the largest weakly connected component in the subgraph.
nx_graph = to_networkx(new_subg_homo.cpu())
largest_cc_nids = list(
max(nx.weakly_connected_components(nx_graph), key=len)
)
largest_cc_homo = node_subgraph(new_subg_homo, largest_cc_nids)
largest_cc_hetero = to_heterogeneous(
largest_cc_homo, new_subg.ntypes, new_subg.etypes
)
# Follow steps for backtracking to original graph node ids
# 1. retrieve instanced homograph from connected-component homograph
# 2. retrieve instanced heterograph from instanced homograph
# 3. retrieve hetero-subgraph from instanced heterograph
# 4. retrieve orignal graph ids from subgraph node ids
cc_nodes = {
ntype: subg.ndata[NID][ntype][
new_subg.ndata[NID][ntype][
new_subg_homo.ndata[NID][
largest_cc_homo.ndata[NID][indicies]
]
]
]
for ntype, indicies in largest_cc_hetero.ndata[NID].items()
}
else:
available_ntypes = [
ntype
for ntype in new_subg.ntypes
if new_subg.num_nodes(ntype) > 0
]
chosen_ntype = np.random.choice(available_ntypes)
# backtrack from subgraph node ids to entire graph
chosen_node = subg.ndata[NID][chosen_ntype][
np.random.choice(new_subg.nodes[chosen_ntype].data[NID])
]
cc_nodes = {
chosen_ntype: torch.tensor(
[chosen_node],
device=subg.device,
)
}
if str(cc_nodes) not in self.mcts_node_maps:
child_mcts_node = MCTSNode(cc_nodes)
self.mcts_node_maps[str(child_mcts_node)] = child_mcts_node
else:
child_mcts_node = self.mcts_node_maps[str(cc_nodes)]
if str(child_mcts_node) not in mcts_children_maps:
mcts_children_maps[str(child_mcts_node)] = child_mcts_node
mcts_node.children = list(mcts_children_maps.values())
for child_mcts_node in mcts_node.children:
if child_mcts_node.immediate_reward == 0:
child_mcts_node.immediate_reward = self.shapley(
child_mcts_node.nodes
)
return mcts_node.children
def mcts_rollout(self, mcts_node):
r"""Perform a MCTS rollout.
Parameters
----------
mcts_node : MCTSNode
Starting node for MCTS
Returns
-------
float
Reward for visiting the node this time
"""
if (
sum(len(nodes) for nodes in mcts_node.nodes.values())
<= self.node_min
):
return mcts_node.immediate_reward
children_nodes = self.get_mcts_children(mcts_node)
children_visit_sum = sum([child.num_visit for child in children_nodes])
children_visit_sum_sqrt = math.sqrt(children_visit_sum)
chosen_child = max(
children_nodes,
key=lambda c: c.total_reward / max(c.num_visit, 1)
+ self.coef
* c.immediate_reward
* children_visit_sum_sqrt
/ (1 + c.num_visit),
)
reward = self.mcts_rollout(chosen_child)
chosen_child.num_visit += 1
chosen_child.total_reward += reward
return reward
def explain_graph(self, graph, feat, target_class, **kwargs):
r"""Find the most important subgraph from the original graph for the
model to classify the graph into the target class.
Parameters
----------
graph : DGLGraph
A heterogeneous graph
feat : dict[str, Tensor]
The dictionary that associates input node features (values) with
the respective node types (keys) present in the graph.
The input features are of shape :math:`(N_t, D_t)`. :math:`N_t` is the
number of nodes for node type :math:`t`, and :math:`D_t` is the feature size for
node type :math:`t`
target_class : int
The target class to explain
kwargs : dict
Additional arguments passed to the GNN model
Returns
-------
dict[str, Tensor]
The dictionary associating tensor node ids (values) to
node types (keys) that represents the most important subgraph
Examples
--------
>>> import dgl
>>> import dgl.function as fn
>>> import torch as th
>>> import torch.nn as nn
>>> import torch.nn.functional as F
>>> from dgl.nn import HeteroSubgraphX
>>> class Model(nn.Module):
... def __init__(self, in_dim, num_classes, canonical_etypes):
... super(Model, self).__init__()
... self.etype_weights = nn.ModuleDict(
... {
... "_".join(c_etype): nn.Linear(in_dim, num_classes)
... for c_etype in canonical_etypes
... }
... )
...
... def forward(self, graph, feat):
... with graph.local_scope():
... c_etype_func_dict = {}
... for c_etype in graph.canonical_etypes:
... src_type, etype, dst_type = c_etype
... wh = self.etype_weights["_".join(c_etype)](feat[src_type])
... graph.nodes[src_type].data[f"h_{c_etype}"] = wh
... c_etype_func_dict[c_etype] = (
... fn.copy_u(f"h_{c_etype}", "m"),
... fn.mean("m", "h"),
... )
... graph.multi_update_all(c_etype_func_dict, "sum")
... hg = 0
... for ntype in graph.ntypes:
... if graph.num_nodes(ntype):
... hg = hg + dgl.mean_nodes(graph, "h", ntype=ntype)
... return hg
>>> input_dim = 5
>>> num_classes = 2
>>> g = dgl.heterograph({("user", "plays", "game"): ([0, 1, 1, 2], [0, 0, 1, 1])})
>>> g.nodes["user"].data["h"] = th.randn(g.num_nodes("user"), input_dim)
>>> g.nodes["game"].data["h"] = th.randn(g.num_nodes("game"), input_dim)
>>> transform = dgl.transforms.AddReverse()
>>> g = transform(g)
>>> # define and train the model
>>> model = Model(input_dim, num_classes, g.canonical_etypes)
>>> feat = g.ndata["h"]
>>> optimizer = th.optim.Adam(model.parameters())
>>> for epoch in range(10):
... logits = model(g, feat)
... loss = F.cross_entropy(logits, th.tensor([1]))
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Explain for the graph
>>> explainer = HeteroSubgraphX(model, num_hops=1)
>>> explainer.explain_graph(g, feat, target_class=1)
{'game': tensor([0, 1]), 'user': tensor([1, 2])}
"""
self.model.eval()
assert (
graph.num_nodes() > self.node_min
), f"The number of nodes in the\
graph {graph.num_nodes()} should be bigger than {self.node_min}."
self.graph = graph
self.feat = feat
self.target_class = target_class
self.kwargs = kwargs
# book all nodes in MCTS
self.mcts_node_maps = dict()
root_dict = {ntype: graph.nodes(ntype) for ntype in graph.ntypes}
root = MCTSNode(root_dict)
self.mcts_node_maps[str(root)] = root
for i in range(self.num_rollouts):
if self.log:
print(
f"Rollout {i}/{self.num_rollouts}, \
{len(self.mcts_node_maps)} subgraphs have been explored."
)
self.mcts_rollout(root)
best_leaf = None
best_immediate_reward = float("-inf")
for mcts_node in self.mcts_node_maps.values():
if len(mcts_node.nodes) > self.node_min:
continue
if mcts_node.immediate_reward > best_immediate_reward:
best_leaf = mcts_node
best_immediate_reward = best_leaf.immediate_reward
return best_leaf.nodes
python/dgl/nn/pytorch/factory.py
+384
View File
@@ -0,0 +1,384 @@
"""Modules that transforms between graphs and between graph and tensors."""
import torch.nn as nn
from ...transforms import knn_graph, radius_graph, segmented_knn_graph
def pairwise_squared_distance(x):
"""
x : (n_samples, n_points, dims)
return : (n_samples, n_points, n_points)
"""
x2s = (x * x).sum(-1, keepdim=True)
return x2s + x2s.transpose(-1, -2) - 2 * x @ x.transpose(-1, -2)
class KNNGraph(nn.Module):
r"""Layer that transforms one point set into a graph, or a batch of
point sets with the same number of points into a batched union of those graphs.
The KNNGraph is implemented in the following steps:
1. Compute an NxN matrix of pairwise distance for all points.
2. Pick the k points with the smallest distance for each point as their k-nearest neighbors.
3. Construct a graph with edges to each point as a node from its k-nearest neighbors.
The overall computational complexity is :math:`O(N^2(logN + D)`.
If a batch of point sets is provided, the point :math:`j` in point
set :math:`i` is mapped to graph node ID: :math:`i \times M + j`, where
:math:`M` is the number of nodes in each point set.
The predecessors of each node are the k-nearest neighbors of the
corresponding point.
Parameters
----------
k : int
The number of neighbors.
Notes
-----
The nearest neighbors found for a node include the node itself.
Examples
--------
The following example uses PyTorch backend.
>>> import torch
>>> from dgl.nn.pytorch.factory import KNNGraph
>>>
>>> kg = KNNGraph(2)
>>> x = torch.tensor([[0,1],
[1,2],
[1,3],
[100, 101],
[101, 102],
[50, 50]])
>>> g = kg(x)
>>> print(g.edges())
(tensor([0, 1, 1, 1, 2, 2, 2, 3, 3, 4, 4, 5]),
tensor([0, 0, 1, 2, 1, 2, 5, 3, 4, 3, 4, 5]))
"""
def __init__(self, k):
super(KNNGraph, self).__init__()
self.k = k
# pylint: disable=invalid-name
def forward(
self,
x,
algorithm="bruteforce-blas",
dist="euclidean",
exclude_self=False,
):
r"""
Forward computation.
Parameters
----------
x : Tensor
:math:`(M, D)` or :math:`(N, M, D)` where :math:`N` means the
number of point sets, :math:`M` means the number of points in
each point set, and :math:`D` means the size of features.
algorithm : str, optional
Algorithm used to compute the k-nearest neighbors.
* 'bruteforce-blas' will first compute the distance matrix
using BLAS matrix multiplication operation provided by
backend frameworks. Then use topk algorithm to get
k-nearest neighbors. This method is fast when the point
set is small but has :math:`O(N^2)` memory complexity where
:math:`N` is the number of points.
* 'bruteforce' will compute distances pair by pair and
directly select the k-nearest neighbors during distance
computation. This method is slower than 'bruteforce-blas'
but has less memory overhead (i.e., :math:`O(Nk)` where :math:`N`
is the number of points, :math:`k` is the number of nearest
neighbors per node) since we do not need to store all distances.
* 'bruteforce-sharemem' (CUDA only) is similar to 'bruteforce'
but use shared memory in CUDA devices for buffer. This method is
faster than 'bruteforce' when the dimension of input points
is not large. This method is only available on CUDA device.
* 'kd-tree' will use the kd-tree algorithm (CPU only).
This method is suitable for low-dimensional data (e.g. 3D
point clouds)
* 'nn-descent' is a approximate approach from paper
`Efficient k-nearest neighbor graph construction for generic similarity
measures <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_. This method
will search for nearest neighbor candidates in "neighbors' neighbors".
(default: 'bruteforce-blas')
dist : str, optional
The distance metric used to compute distance between points. It can be the following
metrics:
* 'euclidean': Use Euclidean distance (L2 norm)
:math:`\sqrt{\sum_{i} (x_{i} - y_{i})^{2}}`.
* 'cosine': Use cosine distance.
(default: 'euclidean')
exclude_self : bool, optional
If True, the output graph will not contain self loop edges, and each node will not
be counted as one of its own k neighbors. If False, the output graph will contain
self loop edges, and a node will be counted as one of its own k neighbors.
Returns
-------
DGLGraph
A DGLGraph without features.
"""
return knn_graph(
x, self.k, algorithm=algorithm, dist=dist, exclude_self=exclude_self
)
class SegmentedKNNGraph(nn.Module):
r"""Layer that transforms one point set into a graph, or a batch of
point sets with different number of points into a batched union of those graphs.
If a batch of point sets is provided, then the point :math:`j` in the point
set :math:`i` is mapped to graph node ID:
:math:`\sum_{p<i} |V_p| + j`, where :math:`|V_p|` means the number of
points in the point set :math:`p`.
The predecessors of each node are the k-nearest neighbors of the
corresponding point.
Parameters
----------
k : int
The number of neighbors.
Notes
-----
The nearest neighbors found for a node include the node itself.
Examples
--------
The following example uses PyTorch backend.
>>> import torch
>>> from dgl.nn.pytorch.factory import SegmentedKNNGraph
>>>
>>> kg = SegmentedKNNGraph(2)
>>> x = torch.tensor([[0,1],
... [1,2],
... [1,3],
... [100, 101],
... [101, 102],
... [50, 50],
... [24,25],
... [25,24]])
>>> g = kg(x, [3,3,2])
>>> print(g.edges())
(tensor([0, 1, 1, 1, 2, 2, 3, 3, 3, 4, 4, 5, 6, 6, 7, 7]),
tensor([0, 0, 1, 2, 1, 2, 3, 4, 5, 3, 4, 5, 6, 7, 6, 7]))
>>>
"""
def __init__(self, k):
super(SegmentedKNNGraph, self).__init__()
self.k = k
# pylint: disable=invalid-name
def forward(
self,
x,
segs,
algorithm="bruteforce-blas",
dist="euclidean",
exclude_self=False,
):
r"""Forward computation.
Parameters
----------
x : Tensor
:math:`(M, D)` where :math:`M` means the total number of points
in all point sets, and :math:`D` means the size of features.
segs : iterable of int
:math:`(N)` integers where :math:`N` means the number of point
sets. The number of elements must sum up to :math:`M`. And any
:math:`N` should :math:`\ge k`
algorithm : str, optional
Algorithm used to compute the k-nearest neighbors.
* 'bruteforce-blas' will first compute the distance matrix
using BLAS matrix multiplication operation provided by
backend frameworks. Then use topk algorithm to get
k-nearest neighbors. This method is fast when the point
set is small but has :math:`O(N^2)` memory complexity where
:math:`N` is the number of points.
* 'bruteforce' will compute distances pair by pair and
directly select the k-nearest neighbors during distance
computation. This method is slower than 'bruteforce-blas'
but has less memory overhead (i.e., :math:`O(Nk)` where :math:`N`
is the number of points, :math:`k` is the number of nearest
neighbors per node) since we do not need to store all distances.
* 'bruteforce-sharemem' (CUDA only) is similar to 'bruteforce'
but use shared memory in CUDA devices for buffer. This method is
faster than 'bruteforce' when the dimension of input points
is not large. This method is only available on CUDA device.
* 'kd-tree' will use the kd-tree algorithm (CPU only).
This method is suitable for low-dimensional data (e.g. 3D
point clouds)
* 'nn-descent' is a approximate approach from paper
`Efficient k-nearest neighbor graph construction for generic similarity
measures <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_. This method
will search for nearest neighbor candidates in "neighbors' neighbors".
(default: 'bruteforce-blas')
dist : str, optional
The distance metric used to compute distance between points. It can be the following
metrics:
* 'euclidean': Use Euclidean distance (L2 norm)
:math:`\sqrt{\sum_{i} (x_{i} - y_{i})^{2}}`.
* 'cosine': Use cosine distance.
(default: 'euclidean')
exclude_self : bool, optional
If True, the output graph will not contain self loop edges, and each node will not
be counted as one of its own k neighbors. If False, the output graph will contain
self loop edges, and a node will be counted as one of its own k neighbors.
Returns
-------
DGLGraph
A batched DGLGraph without features.
"""
return segmented_knn_graph(
x,
self.k,
segs,
algorithm=algorithm,
dist=dist,
exclude_self=exclude_self,
)
class RadiusGraph(nn.Module):
r"""Layer that transforms one point set into a bidirected graph with
neighbors within given distance.
The RadiusGraph is implemented in the following steps:
1. Compute an NxN matrix of pairwise distance for all points.
2. Pick the points within distance to each point as their neighbors.
3. Construct a graph with edges to each point as a node from its neighbors.
The nodes of the returned graph correspond to the points, where the neighbors
of each point are within given distance.
Parameters
----------
r : float
Radius of the neighbors.
p : float, optional
Power parameter for the Minkowski metric. When :attr:`p = 1` it is the
equivalent of Manhattan distance (L1 norm) and Euclidean distance
(L2 norm) for :attr:`p = 2`.
(default: 2)
self_loop : bool, optional
Whether the radius graph will contain self-loops.
(default: False)
compute_mode : str, optional
``use_mm_for_euclid_dist_if_necessary`` - will use matrix multiplication
approach to calculate euclidean distance (p = 2) if P > 25 or R > 25
``use_mm_for_euclid_dist`` - will always use matrix multiplication
approach to calculate euclidean distance (p = 2)
``donot_use_mm_for_euclid_dist`` - will never use matrix multiplication
approach to calculate euclidean distance (p = 2).
(default: donot_use_mm_for_euclid_dist)
Examples
--------
The following examples uses PyTorch backend.
>>> import dgl
>>> from dgl.nn.pytorch.factory import RadiusGraph
>>> x = torch.tensor([[0.0, 0.0, 1.0],
... [1.0, 0.5, 0.5],
... [0.5, 0.2, 0.2],
... [0.3, 0.2, 0.4]])
>>> rg = RadiusGraph(0.75)
>>> g = rg(x) # Each node has neighbors within 0.75 distance
>>> g.edges()
(tensor([0, 1, 2, 2, 3, 3]), tensor([3, 2, 1, 3, 0, 2]))
When :attr:`get_distances` is True, forward pass returns the radius graph and
distances for the corresponding edges.
>>> x = torch.tensor([[0.0, 0.0, 1.0],
... [1.0, 0.5, 0.5],
... [0.5, 0.2, 0.2],
... [0.3, 0.2, 0.4]])
>>> rg = RadiusGraph(0.75)
>>> g, dist = rg(x, get_distances=True)
>>> g.edges()
(tensor([0, 1, 2, 2, 3, 3]), tensor([3, 2, 1, 3, 0, 2]))
>>> dist
tensor([[0.7000],
[0.6557],
[0.6557],
[0.2828],
[0.7000],
[0.2828]])
"""
# pylint: disable=invalid-name
def __init__(
self,
r,
p=2,
self_loop=False,
compute_mode="donot_use_mm_for_euclid_dist",
):
super(RadiusGraph, self).__init__()
self.r = r
self.p = p
self.self_loop = self_loop
self.compute_mode = compute_mode
# pylint: disable=invalid-name
def forward(self, x, get_distances=False):
r"""
Forward computation.
Parameters
----------
x : Tensor
The point coordinates. :math:`(N, D)` where :math:`N` means the
number of points in the point set, and :math:`D` means the size of
the features. It can be either on CPU or GPU. Device of the point
coordinates specifies device of the radius graph.
get_distances : bool, optional
Whether to return the distances for the corresponding edges in the
radius graph.
(default: False)
Returns
-------
DGLGraph
The constructed graph. The node IDs are in the same order as :attr:`x`.
torch.Tensor, optional
The distances for the edges in the constructed graph. The distances
are in the same order as edge IDs.
"""
return radius_graph(
x, self.r, self.p, self.self_loop, self.compute_mode, get_distances
)
python/dgl/nn/pytorch/glob.py
+1393
View File
File diff suppressed because it is too large Load Diff
python/dgl/nn/pytorch/gt/__init__.py
+9
View File
@@ -0,0 +1,9 @@
"""Torch modules for Graph Transformer."""
from .biased_mha import BiasedMHA
from .degree_encoder import DegreeEncoder
from .egt import EGTLayer
from .graphormer import GraphormerLayer
from .lap_pos_encoder import LapPosEncoder
from .path_encoder import PathEncoder
from .spatial_encoder import SpatialEncoder, SpatialEncoder3d
python/dgl/nn/pytorch/gt/biased_mha.py
+158
View File
@@ -0,0 +1,158 @@
"""Biased Multi-head Attention"""
import torch as th
import torch.nn as nn
import torch.nn.functional as F
class BiasedMHA(nn.Module):
r"""Dense Multi-Head Attention Module with Graph Attention Bias.
Compute attention between nodes with attention bias obtained from graph
structures, as introduced in `Do Transformers Really Perform Bad for
Graph Representation? <https://arxiv.org/pdf/2106.05234>`__
.. math::
\text{Attn}=\text{softmax}(\dfrac{QK^T}{\sqrt{d}} \circ b)
:math:`Q` and :math:`K` are feature representations of nodes. :math:`d`
is the corresponding :attr:`feat_size`. :math:`b` is attention bias, which
can be additive or multiplicative according to the operator :math:`\circ`.
Parameters
----------
feat_size : int
Feature size.
num_heads : int
Number of attention heads, by which :attr:`feat_size` is divisible.
bias : bool, optional
If True, it uses bias for linear projection. Default: True.
attn_bias_type : str, optional
The type of attention bias used for modifying attention. Selected from
'add' or 'mul'. Default: 'add'.
* 'add' is for additive attention bias.
* 'mul' is for multiplicative attention bias.
attn_drop : float, optional
Dropout probability on attention weights. Defalt: 0.1.
Examples
--------
>>> import torch as th
>>> from dgl.nn import BiasedMHA
>>> ndata = th.rand(16, 100, 512)
>>> bias = th.rand(16, 100, 100, 8)
>>> net = BiasedMHA(feat_size=512, num_heads=8)
>>> out = net(ndata, bias)
"""
def __init__(
self,
feat_size,
num_heads,
bias=True,
attn_bias_type="add",
attn_drop=0.1,
):
super().__init__()
self.feat_size = feat_size
self.num_heads = num_heads
self.head_dim = feat_size // num_heads
assert (
self.head_dim * num_heads == feat_size
), "feat_size must be divisible by num_heads"
self.scaling = self.head_dim**-0.5
self.attn_bias_type = attn_bias_type
self.q_proj = nn.Linear(feat_size, feat_size, bias=bias)
self.k_proj = nn.Linear(feat_size, feat_size, bias=bias)
self.v_proj = nn.Linear(feat_size, feat_size, bias=bias)
self.out_proj = nn.Linear(feat_size, feat_size, bias=bias)
self.dropout = nn.Dropout(p=attn_drop)
self.reset_parameters()
def reset_parameters(self):
"""
Initialize parameters of projection matrices, the same settings as in
the original implementation of the paper.
"""
nn.init.xavier_uniform_(self.q_proj.weight, gain=2**-0.5)
nn.init.xavier_uniform_(self.k_proj.weight, gain=2**-0.5)
nn.init.xavier_uniform_(self.v_proj.weight, gain=2**-0.5)
nn.init.xavier_uniform_(self.out_proj.weight)
if self.out_proj.bias is not None:
nn.init.constant_(self.out_proj.bias, 0.0)
def forward(self, ndata, attn_bias=None, attn_mask=None):
"""Forward computation.
Parameters
----------
ndata : torch.Tensor
A 3D input tensor. Shape: (batch_size, N, :attr:`feat_size`), where
N is the maximum number of nodes.
attn_bias : torch.Tensor, optional
The attention bias used for attention modification. Shape:
(batch_size, N, N, :attr:`num_heads`).
attn_mask : torch.Tensor, optional
The attention mask used for avoiding computation on invalid
positions, where invalid positions are indicated by `True` values.
Shape: (batch_size, N, N). Note: For rows corresponding to
unexisting nodes, make sure at least one entry is set to `False` to
prevent obtaining NaNs with softmax.
Returns
-------
y : torch.Tensor
The output tensor. Shape: (batch_size, N, :attr:`feat_size`)
"""
q_h = self.q_proj(ndata).transpose(0, 1)
k_h = self.k_proj(ndata).transpose(0, 1)
v_h = self.v_proj(ndata).transpose(0, 1)
bsz, N, _ = ndata.shape
q_h = (
q_h.reshape(N, bsz * self.num_heads, self.head_dim).transpose(0, 1)
* self.scaling
)
k_h = k_h.reshape(N, bsz * self.num_heads, self.head_dim).permute(
1, 2, 0
)
v_h = v_h.reshape(N, bsz * self.num_heads, self.head_dim).transpose(
0, 1
)
attn_weights = (
th.bmm(q_h, k_h)
.transpose(0, 2)
.reshape(N, N, bsz, self.num_heads)
.transpose(0, 2)
)
if attn_bias is not None:
if self.attn_bias_type == "add":
attn_weights += attn_bias
else:
attn_weights *= attn_bias
if attn_mask is not None:
attn_weights[attn_mask.to(th.bool)] = float("-inf")
attn_weights = F.softmax(
attn_weights.transpose(0, 2)
.reshape(N, N, bsz * self.num_heads)
.transpose(0, 2),
dim=2,
)
attn_weights = self.dropout(attn_weights)
attn = th.bmm(attn_weights, v_h).transpose(0, 1)
attn = self.out_proj(
attn.reshape(N, bsz, self.feat_size).transpose(0, 1)
)
return attn
python/dgl/nn/pytorch/gt/degree_encoder.py
+98
View File
@@ -0,0 +1,98 @@
"""Degree Encoder"""
import torch as th
import torch.nn as nn
class DegreeEncoder(nn.Module):
r"""Degree Encoder, as introduced in
`Do Transformers Really Perform Bad for Graph Representation?
<https://proceedings.neurips.cc/paper/2021/file/f1c1592588411002af340cbaedd6fc33-Paper.pdf>`__
This module is a learnable degree embedding module.
Parameters
----------
max_degree : int
Upper bound of degrees to be encoded.
Each degree will be clamped into the range [0, ``max_degree``].
embedding_dim : int
Output dimension of embedding vectors.
direction : str, optional
Degrees of which direction to be encoded,
selected from ``in``, ``out`` and ``both``.
``both`` encodes degrees from both directions
and output the addition of them.
Default : ``both``.
Example
-------
>>> import dgl
>>> from dgl.nn import DegreeEncoder
>>> import torch as th
>>> from torch.nn.utils.rnn import pad_sequence
>>> g1 = dgl.graph(([0,0,0,1,1,2,3,3], [1,2,3,0,3,0,0,1]))
>>> g2 = dgl.graph(([0,1], [1,0]))
>>> in_degree = pad_sequence([g1.in_degrees(), g2.in_degrees()], batch_first=True)
>>> out_degree = pad_sequence([g1.out_degrees(), g2.out_degrees()], batch_first=True)
>>> print(in_degree.shape)
torch.Size([2, 4])
>>> degree_encoder = DegreeEncoder(5, 16)
>>> degree_embedding = degree_encoder(th.stack((in_degree, out_degree)))
>>> print(degree_embedding.shape)
torch.Size([2, 4, 16])
"""
def __init__(self, max_degree, embedding_dim, direction="both"):
super(DegreeEncoder, self).__init__()
self.direction = direction
if direction == "both":
self.encoder1 = nn.Embedding(
max_degree + 1, embedding_dim, padding_idx=0
)
self.encoder2 = nn.Embedding(
max_degree + 1, embedding_dim, padding_idx=0
)
else:
self.encoder = nn.Embedding(
max_degree + 1, embedding_dim, padding_idx=0
)
self.max_degree = max_degree
def forward(self, degrees):
"""
Parameters
----------
degrees : Tensor
If :attr:`direction` is ``both``, it should be stacked in degrees and out degrees
of the batched graph with zero padding, a tensor of shape :math:`(2, B, N)`.
Otherwise, it should be zero-padded in degrees or out degrees of the batched
graph, a tensor of shape :math:`(B, N)`, where :math:`B` is the batch size
of the batched graph, and :math:`N` is the maximum number of nodes.
Returns
-------
Tensor
Return degree embedding vectors of shape :math:`(B, N, d)`,
where :math:`d` is :attr:`embedding_dim`.
"""
degrees = th.clamp(degrees, min=0, max=self.max_degree)
if self.direction == "in":
assert len(degrees.shape) == 2
degree_embedding = self.encoder(degrees)
elif self.direction == "out":
assert len(degrees.shape) == 2
degree_embedding = self.encoder(degrees)
elif self.direction == "both":
assert len(degrees.shape) == 3 and degrees.shape[0] == 2
degree_embedding = self.encoder1(degrees[0]) + self.encoder2(
degrees[1]
)
else:
raise ValueError(
f'Supported direction options: "in", "out" and "both", '
f"but got {self.direction}"
)
return degree_embedding
python/dgl/nn/pytorch/gt/egt.py
+177
View File
@@ -0,0 +1,177 @@
"""EGT Layer"""
import torch
import torch.nn as nn
import torch.nn.functional as F
class EGTLayer(nn.Module):
r"""EGTLayer for Edge-augmented Graph Transformer (EGT), as introduced in
`Global Self-Attention as a Replacement for Graph Convolution
Reference `<https://arxiv.org/pdf/2108.03348.pdf>`_
Parameters
----------
feat_size : int
Node feature size.
edge_feat_size : int
Edge feature size.
num_heads : int
Number of attention heads, by which :attr: `feat_size` is divisible.
num_virtual_nodes : int
Number of virtual nodes.
dropout : float, optional
Dropout probability. Default: 0.0.
attn_dropout : float, optional
Attention dropout probability. Default: 0.0.
activation : callable activation layer, optional
Activation function. Default: nn.ELU().
edge_update : bool, optional
Whether to update the edge embedding. Default: True.
Examples
--------
>>> import torch as th
>>> from dgl.nn import EGTLayer
>>> batch_size = 16
>>> num_nodes = 100
>>> feat_size, edge_feat_size = 128, 32
>>> nfeat = th.rand(batch_size, num_nodes, feat_size)
>>> efeat = th.rand(batch_size, num_nodes, num_nodes, edge_feat_size)
>>> net = EGTLayer(
feat_size=feat_size,
edge_feat_size=edge_feat_size,
num_heads=8,
num_virtual_nodes=4,
)
>>> out = net(nfeat, efeat)
"""
def __init__(
self,
feat_size,
edge_feat_size,
num_heads,
num_virtual_nodes,
dropout=0,
attn_dropout=0,
activation=nn.ELU(),
edge_update=True,
):
super().__init__()
self.num_heads = num_heads
self.num_virtual_nodes = num_virtual_nodes
self.edge_update = edge_update
assert (
feat_size % num_heads == 0
), "feat_size must be divisible by num_heads"
self.dot_dim = feat_size // num_heads
self.mha_ln_h = nn.LayerNorm(feat_size)
self.mha_ln_e = nn.LayerNorm(edge_feat_size)
self.edge_input = nn.Linear(edge_feat_size, num_heads)
self.qkv_proj = nn.Linear(feat_size, feat_size * 3)
self.gate = nn.Linear(edge_feat_size, num_heads)
self.attn_dropout = nn.Dropout(attn_dropout)
self.node_output = nn.Linear(feat_size, feat_size)
self.mha_dropout_h = nn.Dropout(dropout)
self.node_ffn = nn.Sequential(
nn.LayerNorm(feat_size),
nn.Linear(feat_size, feat_size),
activation,
nn.Linear(feat_size, feat_size),
nn.Dropout(dropout),
)
if self.edge_update:
self.edge_output = nn.Linear(num_heads, edge_feat_size)
self.mha_dropout_e = nn.Dropout(dropout)
self.edge_ffn = nn.Sequential(
nn.LayerNorm(edge_feat_size),
nn.Linear(edge_feat_size, edge_feat_size),
activation,
nn.Linear(edge_feat_size, edge_feat_size),
nn.Dropout(dropout),
)
def forward(self, nfeat, efeat, mask=None):
"""Forward computation. Note: :attr:`nfeat` and :attr:`efeat` should be
padded with embedding of virtual nodes if :attr:`num_virtual_nodes` > 0,
while :attr:`mask` should be padded with `0` values for virtual nodes.
The padding should be put at the beginning.
Parameters
----------
nfeat : torch.Tensor
A 3D input tensor. Shape: (batch_size, N, :attr:`feat_size`), where N
is the sum of the maximum number of nodes and the number of virtual nodes.
efeat : torch.Tensor
Edge embedding used for attention computation and self update.
Shape: (batch_size, N, N, :attr:`edge_feat_size`).
mask : torch.Tensor, optional
The attention mask used for avoiding computation on invalid
positions, where valid positions are indicated by `0` and
invalid positions are indicated by `-inf`.
Shape: (batch_size, N, N). Default: None.
Returns
-------
nfeat : torch.Tensor
The output node embedding. Shape: (batch_size, N, :attr:`feat_size`).
efeat : torch.Tensor, optional
The output edge embedding. Shape: (batch_size, N, N, :attr:`edge_feat_size`).
It is returned only if :attr:`edge_update` is True.
"""
nfeat_r1 = nfeat
efeat_r1 = efeat
nfeat_ln = self.mha_ln_h(nfeat)
efeat_ln = self.mha_ln_e(efeat)
qkv = self.qkv_proj(nfeat_ln)
e_bias = self.edge_input(efeat_ln)
gates = self.gate(efeat_ln)
bsz, N, _ = qkv.shape
q_h, k_h, v_h = qkv.view(bsz, N, -1, self.num_heads).split(
self.dot_dim, dim=2
)
attn_hat = torch.einsum("bldh,bmdh->blmh", q_h, k_h)
attn_hat = attn_hat.clamp(-5, 5) + e_bias
if mask is None:
gates = torch.sigmoid(gates)
attn_tild = F.softmax(attn_hat, dim=2) * gates
else:
gates = torch.sigmoid(gates + mask.unsqueeze(-1))
attn_tild = F.softmax(attn_hat + mask.unsqueeze(-1), dim=2) * gates
attn_tild = self.attn_dropout(attn_tild)
v_attn = torch.einsum("blmh,bmkh->blkh", attn_tild, v_h)
# Scale the aggregated values by degree.
degrees = torch.sum(gates, dim=2, keepdim=True)
degree_scalers = torch.log(1 + degrees)
degree_scalers[:, : self.num_virtual_nodes] = 1.0
v_attn = v_attn * degree_scalers
v_attn = v_attn.reshape(bsz, N, self.num_heads * self.dot_dim)
nfeat = self.node_output(v_attn)
nfeat = self.mha_dropout_h(nfeat)
nfeat.add_(nfeat_r1)
nfeat_r2 = nfeat
nfeat = self.node_ffn(nfeat)
nfeat.add_(nfeat_r2)
if self.edge_update:
efeat = self.edge_output(attn_hat)
efeat = self.mha_dropout_e(efeat)
efeat.add_(efeat_r1)
efeat_r2 = efeat
efeat = self.edge_ffn(efeat)
efeat.add_(efeat_r2)
return nfeat, efeat
return nfeat
python/dgl/nn/pytorch/gt/graphormer.py
+128
View File
@@ -0,0 +1,128 @@
"""Graphormer Layer"""
import torch.nn as nn
from .biased_mha import BiasedMHA
class GraphormerLayer(nn.Module):
r"""Graphormer Layer with Dense Multi-Head Attention, as introduced
in `Do Transformers Really Perform Bad for Graph Representation?
<https://arxiv.org/pdf/2106.05234>`__
Parameters
----------
feat_size : int
Feature size.
hidden_size : int
Hidden size of feedforward layers.
num_heads : int
Number of attention heads, by which :attr:`feat_size` is divisible.
attn_bias_type : str, optional
The type of attention bias used for modifying attention. Selected from
'add' or 'mul'. Default: 'add'.
* 'add' is for additive attention bias.
* 'mul' is for multiplicative attention bias.
norm_first : bool, optional
If True, it performs layer normalization before attention and
feedforward operations. Otherwise, it applies layer normalization
afterwards. Default: False.
dropout : float, optional
Dropout probability. Default: 0.1.
attn_dropout : float, optional
Attention dropout probability. Default: 0.1.
activation : callable activation layer, optional
Activation function. Default: nn.ReLU().
Examples
--------
>>> import torch as th
>>> from dgl.nn import GraphormerLayer
>>> batch_size = 16
>>> num_nodes = 100
>>> feat_size = 512
>>> num_heads = 8
>>> nfeat = th.rand(batch_size, num_nodes, feat_size)
>>> bias = th.rand(batch_size, num_nodes, num_nodes, num_heads)
>>> net = GraphormerLayer(
feat_size=feat_size,
hidden_size=2048,
num_heads=num_heads
)
>>> out = net(nfeat, bias)
"""
def __init__(
self,
feat_size,
hidden_size,
num_heads,
attn_bias_type="add",
norm_first=False,
dropout=0.1,
attn_dropout=0.1,
activation=nn.ReLU(),
):
super().__init__()
self.norm_first = norm_first
self.attn = BiasedMHA(
feat_size=feat_size,
num_heads=num_heads,
attn_bias_type=attn_bias_type,
attn_drop=attn_dropout,
)
self.ffn = nn.Sequential(
nn.Linear(feat_size, hidden_size),
activation,
nn.Dropout(p=dropout),
nn.Linear(hidden_size, feat_size),
nn.Dropout(p=dropout),
)
self.dropout = nn.Dropout(p=dropout)
self.attn_layer_norm = nn.LayerNorm(feat_size)
self.ffn_layer_norm = nn.LayerNorm(feat_size)
def forward(self, nfeat, attn_bias=None, attn_mask=None):
"""Forward computation.
Parameters
----------
nfeat : torch.Tensor
A 3D input tensor. Shape: (batch_size, N, :attr:`feat_size`), where
N is the maximum number of nodes.
attn_bias : torch.Tensor, optional
The attention bias used for attention modification. Shape:
(batch_size, N, N, :attr:`num_heads`).
attn_mask : torch.Tensor, optional
The attention mask used for avoiding computation on invalid
positions, where invalid positions are indicated by `True` values.
Shape: (batch_size, N, N). Note: For rows corresponding to
unexisting nodes, make sure at least one entry is set to `False` to
prevent obtaining NaNs with softmax.
Returns
-------
y : torch.Tensor
The output tensor. Shape: (batch_size, N, :attr:`feat_size`)
"""
residual = nfeat
if self.norm_first:
nfeat = self.attn_layer_norm(nfeat)
nfeat = self.attn(nfeat, attn_bias, attn_mask)
nfeat = self.dropout(nfeat)
nfeat = residual + nfeat
if not self.norm_first:
nfeat = self.attn_layer_norm(nfeat)
residual = nfeat
if self.norm_first:
nfeat = self.ffn_layer_norm(nfeat)
nfeat = self.ffn(nfeat)
nfeat = residual + nfeat
if not self.norm_first:
nfeat = self.ffn_layer_norm(nfeat)
return nfeat
python/dgl/nn/pytorch/gt/lap_pos_encoder.py
+162
View File
@@ -0,0 +1,162 @@
"""Laplacian Positional Encoder"""
import torch as th
import torch.nn as nn
class LapPosEncoder(nn.Module):
r"""Laplacian Positional Encoder (LPE), as introduced in
`GraphGPS: General Powerful Scalable Graph Transformers
<https://arxiv.org/abs/2205.12454>`__
This module is a learned laplacian positional encoding module using
Transformer or DeepSet.
Parameters
----------
model_type : str
Encoder model type for LPE, can only be "Transformer" or "DeepSet".
num_layer : int
Number of layers in Transformer/DeepSet Encoder.
k : int
Number of smallest non-trivial eigenvectors.
dim : int
Output size of final laplacian encoding.
n_head : int, optional
Number of heads in Transformer Encoder.
Default : 1.
batch_norm : bool, optional
If True, apply batch normalization on raw laplacian positional
encoding. Default : False.
num_post_layer : int, optional
If num_post_layer > 0, apply an MLP of ``num_post_layer`` layers after
pooling. Default : 0.
Example
-------
>>> import dgl
>>> from dgl import LapPE
>>> from dgl.nn import LapPosEncoder
>>> transform = LapPE(k=5, feat_name='eigvec', eigval_name='eigval', padding=True)
>>> g = dgl.graph(([0,1,2,3,4,2,3,1,4,0], [2,3,1,4,0,0,1,2,3,4]))
>>> g = transform(g)
>>> eigvals, eigvecs = g.ndata['eigval'], g.ndata['eigvec']
>>> transformer_encoder = LapPosEncoder(
model_type="Transformer", num_layer=3, k=5, dim=16, n_head=4
)
>>> pos_encoding = transformer_encoder(eigvals, eigvecs)
>>> deepset_encoder = LapPosEncoder(
model_type="DeepSet", num_layer=3, k=5, dim=16, num_post_layer=2
)
>>> pos_encoding = deepset_encoder(eigvals, eigvecs)
"""
def __init__(
self,
model_type,
num_layer,
k,
dim,
n_head=1,
batch_norm=False,
num_post_layer=0,
):
super(LapPosEncoder, self).__init__()
self.model_type = model_type
self.linear = nn.Linear(2, dim)
if self.model_type == "Transformer":
encoder_layer = nn.TransformerEncoderLayer(
d_model=dim, nhead=n_head, batch_first=True
)
self.pe_encoder = nn.TransformerEncoder(
encoder_layer, num_layers=num_layer
)
elif self.model_type == "DeepSet":
layers = []
if num_layer == 1:
layers.append(nn.ReLU())
else:
self.linear = nn.Linear(2, 2 * dim)
layers.append(nn.ReLU())
for _ in range(num_layer - 2):
layers.append(nn.Linear(2 * dim, 2 * dim))
layers.append(nn.ReLU())
layers.append(nn.Linear(2 * dim, dim))
layers.append(nn.ReLU())
self.pe_encoder = nn.Sequential(*layers)
else:
raise ValueError(
f"model_type '{model_type}' is not allowed, must be "
"'Transformer' or 'DeepSet'."
)
if batch_norm:
self.raw_norm = nn.BatchNorm1d(k)
else:
self.raw_norm = None
if num_post_layer > 0:
layers = []
if num_post_layer == 1:
layers.append(nn.Linear(dim, dim))
layers.append(nn.ReLU())
else:
layers.append(nn.Linear(dim, 2 * dim))
layers.append(nn.ReLU())
for _ in range(num_post_layer - 2):
layers.append(nn.Linear(2 * dim, 2 * dim))
layers.append(nn.ReLU())
layers.append(nn.Linear(2 * dim, dim))
layers.append(nn.ReLU())
self.post_mlp = nn.Sequential(*layers)
else:
self.post_mlp = None
def forward(self, eigvals, eigvecs):
r"""
Parameters
----------
eigvals : Tensor
Laplacian Eigenvalues of shape :math:`(N, k)`, k different
eigenvalues repeat N times, can be obtained by using `LaplacianPE`.
eigvecs : Tensor
Laplacian Eigenvectors of shape :math:`(N, k)`, can be obtained by
using `LaplacianPE`.
Returns
-------
Tensor
Return the laplacian positional encodings of shape :math:`(N, d)`,
where :math:`N` is the number of nodes in the input graph,
:math:`d` is :attr:`dim`.
"""
pos_encoding = th.cat(
(eigvecs.unsqueeze(2), eigvals.unsqueeze(2)), dim=2
).float()
empty_mask = th.isnan(pos_encoding)
pos_encoding[empty_mask] = 0
if self.raw_norm:
pos_encoding = self.raw_norm(pos_encoding)
pos_encoding = self.linear(pos_encoding)
if self.model_type == "Transformer":
pos_encoding = self.pe_encoder(
src=pos_encoding, src_key_padding_mask=empty_mask[:, :, 1]
)
else:
pos_encoding = self.pe_encoder(pos_encoding)
# Remove masked sequences.
pos_encoding[empty_mask[:, :, 1]] = 0
# Sum pooling.
pos_encoding = th.sum(pos_encoding, 1, keepdim=False)
# MLP post pooling.
if self.post_mlp:
pos_encoding = self.post_mlp(pos_encoding)
return pos_encoding
python/dgl/nn/pytorch/gt/path_encoder.py
+86
View File
@@ -0,0 +1,86 @@
"""Path Encoder"""
import torch as th
import torch.nn as nn
class PathEncoder(nn.Module):
r"""Path Encoder, as introduced in Edge Encoding of
`Do Transformers Really Perform Bad for Graph Representation?
<https://proceedings.neurips.cc/paper/2021/file/f1c1592588411002af340cbaedd6fc33-Paper.pdf>`__
This module is a learnable path embedding module and encodes the shortest
path between each pair of nodes as attention bias.
Parameters
----------
max_len : int
Maximum number of edges in each path to be encoded.
Exceeding part of each path will be truncated, i.e.
truncating edges with serial number no less than :attr:`max_len`.
feat_dim : int
Dimension of edge features in the input graph.
num_heads : int, optional
Number of attention heads if multi-head attention mechanism is applied.
Default : 1.
Examples
--------
>>> import torch as th
>>> import dgl
>>> from dgl.nn import PathEncoder
>>> from dgl import shortest_dist
>>> g = dgl.graph(([0,0,0,1,1,2,3,3], [1,2,3,0,3,0,0,1]))
>>> edata = th.rand(8, 16)
>>> # Since shortest_dist returns -1 for unreachable node pairs,
>>> # edata[-1] should be filled with zero padding.
>>> edata = th.cat(
(edata, th.zeros(1, 16)), dim=0
)
>>> dist, path = shortest_dist(g, root=None, return_paths=True)
>>> path_data = edata[path[:, :, :2]]
>>> path_encoder = PathEncoder(2, 16, num_heads=8)
>>> out = path_encoder(dist.unsqueeze(0), path_data.unsqueeze(0))
>>> print(out.shape)
torch.Size([1, 4, 4, 8])
"""
def __init__(self, max_len, feat_dim, num_heads=1):
super().__init__()
self.max_len = max_len
self.feat_dim = feat_dim
self.num_heads = num_heads
self.embedding_table = nn.Embedding(max_len * num_heads, feat_dim)
def forward(self, dist, path_data):
"""
Parameters
----------
dist : Tensor
Shortest path distance matrix of the batched graph with zero padding,
of shape :math:`(B, N, N)`, where :math:`B` is the batch size of
the batched graph, and :math:`N` is the maximum number of nodes.
path_data : Tensor
Edge feature along the shortest path with zero padding, of shape
:math:`(B, N, N, L, d)`, where :math:`L` is the maximum length of
the shortest paths, and :math:`d` is :attr:`feat_dim`.
Returns
-------
torch.Tensor
Return attention bias as path encoding, of shape
:math:`(B, N, N, H)`, where :math:`B` is the batch size of
the input graph, :math:`N` is the maximum number of nodes, and
:math:`H` is :attr:`num_heads`.
"""
shortest_distance = th.clamp(dist, min=1, max=self.max_len)
edge_embedding = self.embedding_table.weight.reshape(
self.max_len, self.num_heads, -1
)
path_encoding = th.div(
th.einsum("bxyld,lhd->bxyh", path_data, edge_embedding).permute(
3, 0, 1, 2
),
shortest_distance,
).permute(1, 2, 3, 0)
return path_encoding
python/dgl/nn/pytorch/gt/spatial_encoder.py
+209
View File
@@ -0,0 +1,209 @@
"""Spatial Encoder"""
import torch as th
import torch.nn as nn
import torch.nn.functional as F
def gaussian(x, mean, std):
"""compute gaussian basis kernel function"""
const_pi = 3.14159
a = (2 * const_pi) ** 0.5
return th.exp(-0.5 * (((x - mean) / std) ** 2)) / (a * std)
class SpatialEncoder(nn.Module):
r"""Spatial Encoder, as introduced in
`Do Transformers Really Perform Bad for Graph Representation?
<https://proceedings.neurips.cc/paper/2021/file/f1c1592588411002af340cbaedd6fc33-Paper.pdf>`__
This module is a learnable spatial embedding module, which encodes
the shortest distance between each node pair for attention bias.
Parameters
----------
max_dist : int
Upper bound of the shortest path distance
between each node pair to be encoded.
All distance will be clamped into the range `[0, max_dist]`.
num_heads : int, optional
Number of attention heads if multi-head attention mechanism is applied.
Default : 1.
Examples
--------
>>> import torch as th
>>> import dgl
>>> from dgl.nn import SpatialEncoder
>>> from dgl import shortest_dist
>>> g1 = dgl.graph(([0,0,0,1,1,2,3,3], [1,2,3,0,3,0,0,1]))
>>> g2 = dgl.graph(([0,1], [1,0]))
>>> n1, n2 = g1.num_nodes(), g2.num_nodes()
>>> # use -1 padding since shortest_dist returns -1 for unreachable node pairs
>>> dist = -th.ones((2, 4, 4), dtype=th.long)
>>> dist[0, :n1, :n1] = shortest_dist(g1, root=None, return_paths=False)
>>> dist[1, :n2, :n2] = shortest_dist(g2, root=None, return_paths=False)
>>> spatial_encoder = SpatialEncoder(max_dist=2, num_heads=8)
>>> out = spatial_encoder(dist)
>>> print(out.shape)
torch.Size([2, 4, 4, 8])
"""
def __init__(self, max_dist, num_heads=1):
super().__init__()
self.max_dist = max_dist
self.num_heads = num_heads
# deactivate node pair between which the distance is -1
self.embedding_table = nn.Embedding(
max_dist + 2, num_heads, padding_idx=0
)
def forward(self, dist):
"""
Parameters
----------
dist : Tensor
Shortest path distance of the batched graph with -1 padding, a tensor
of shape :math:`(B, N, N)`, where :math:`B` is the batch size of
the batched graph, and :math:`N` is the maximum number of nodes.
Returns
-------
torch.Tensor
Return attention bias as spatial encoding of shape
:math:`(B, N, N, H)`, where :math:`H` is :attr:`num_heads`.
"""
spatial_encoding = self.embedding_table(
th.clamp(
dist,
min=-1,
max=self.max_dist,
)
+ 1
)
return spatial_encoding
class SpatialEncoder3d(nn.Module):
r"""3D Spatial Encoder, as introduced in
`One Transformer Can Understand Both 2D & 3D Molecular Data
<https://arxiv.org/pdf/2210.01765.pdf>`__
This module encodes pair-wise relation between node pair :math:`(i,j)` in
the 3D geometric space, according to the Gaussian Basis Kernel function:
:math:`\psi _{(i,j)} ^k = \frac{1}{\sqrt{2\pi} \lvert \sigma^k \rvert}
\exp{\left ( -\frac{1}{2} \left( \frac{\gamma_{(i,j)} \lvert \lvert r_i -
r_j \rvert \rvert + \beta_{(i,j)} - \mu^k}{\lvert \sigma^k \rvert} \right)
^2 \right)}，k=1,...,K,`
where :math:`K` is the number of Gaussian Basis kernels. :math:`r_i` is the
Cartesian coordinate of node :math:`i`.
:math:`\gamma_{(i,j)}, \beta_{(i,j)}` are learnable scaling factors and
biases determined by node types. :math:`\mu^k, \sigma^k` are learnable
centers and standard deviations of the Gaussian Basis kernels.
Parameters
----------
num_kernels : int
Number of Gaussian Basis Kernels to be applied. Each Gaussian Basis
Kernel contains a learnable kernel center and a learnable standard
deviation.
num_heads : int, optional
Number of attention heads if multi-head attention mechanism is applied.
Default : 1.
max_node_type : int, optional
Maximum number of node types. Each node type has a corresponding
learnable scaling factor and a bias. Default : 100.
Examples
--------
>>> import torch as th
>>> import dgl
>>> from dgl.nn import SpatialEncoder3d
>>> coordinate = th.rand(1, 4, 3)
>>> node_type = th.tensor([[1, 0, 2, 1]])
>>> spatial_encoder = SpatialEncoder3d(num_kernels=4,
... num_heads=8,
... max_node_type=3)
>>> out = spatial_encoder(coordinate, node_type=node_type)
>>> print(out.shape)
torch.Size([1, 4, 4, 8])
"""
def __init__(self, num_kernels, num_heads=1, max_node_type=100):
super().__init__()
self.num_kernels = num_kernels
self.num_heads = num_heads
self.max_node_type = max_node_type
self.means = nn.Parameter(th.empty(num_kernels))
self.stds = nn.Parameter(th.empty(num_kernels))
self.linear_layer_1 = nn.Linear(num_kernels, num_kernels)
self.linear_layer_2 = nn.Linear(num_kernels, num_heads)
# There are 2 * max_node_type + 3 pairs of gamma and beta parameters:
# 1. Parameters at position 0 are for default gamma/beta when no node
# type is given
# 2. Parameters at position 1 to max_node_type+1 are for src node types.
# (position 1 is for padded unexisting nodes)
# 3. Parameters at position max_node_type+2 to 2*max_node_type+2 are
# for tgt node types. (position max_node_type+2 is for padded)
# unexisting nodes)
self.gamma = nn.Embedding(2 * max_node_type + 3, 1, padding_idx=0)
self.beta = nn.Embedding(2 * max_node_type + 3, 1, padding_idx=0)
nn.init.uniform_(self.means, 0, 3)
nn.init.uniform_(self.stds, 0, 3)
nn.init.constant_(self.gamma.weight, 1)
nn.init.constant_(self.beta.weight, 0)
def forward(self, coord, node_type=None):
"""
Parameters
----------
coord : torch.Tensor
3D coordinates of nodes in shape :math:`(B, N, 3)`, where :math:`B`
is the batch size, :math:`N`: is the maximum number of nodes.
node_type : torch.Tensor, optional
Node type ids of nodes. Default : None.
* If specified, :attr:`node_type` should be a tensor in shape
:math:`(B, N,)`. The scaling factors in gaussian kernels of each
pair of nodes are determined by their node types.
* Otherwise, :attr:`node_type` will be set to zeros of the same
shape by default.
Returns
-------
torch.Tensor
Return attention bias as 3D spatial encoding of shape
:math:`(B, N, N, H)`, where :math:`H` is :attr:`num_heads`.
"""
bsz, N = coord.shape[:2]
euc_dist = th.cdist(coord, coord, p=2.0) # shape: [B, n, n]
if node_type is None:
node_type = th.zeros([bsz, N, N, 2], device=coord.device).long()
else:
src_node_type = node_type.unsqueeze(-1).repeat(1, 1, N)
tgt_node_type = node_type.unsqueeze(1).repeat(1, N, 1)
node_type = th.stack(
[src_node_type + 2, tgt_node_type + self.max_node_type + 3],
dim=-1,
) # shape: [B, n, n, 2]
# scaled euclidean distance
gamma = self.gamma(node_type).sum(dim=-2) # shape: [B, n, n, 1]
beta = self.beta(node_type).sum(dim=-2) # shape: [B, n, n, 1]
euc_dist = gamma * euc_dist.unsqueeze(-1) + beta # shape: [B, n, n, 1]
# gaussian basis kernel
euc_dist = euc_dist.expand(-1, -1, -1, self.num_kernels)
gaussian_kernel = gaussian(
euc_dist, self.means, self.stds.abs() + 1e-2
) # shape: [B, n, n, K]
# linear projection
encoding = self.linear_layer_1(gaussian_kernel)
encoding = F.gelu(encoding)
encoding = self.linear_layer_2(encoding) # shape: [B, n, n, H]
return encoding
python/dgl/nn/pytorch/hetero.py
+428
View File
@@ -0,0 +1,428 @@
"""Heterograph NN modules"""
from functools import partial
import torch as th
import torch.nn as nn
from ...base import DGLError
__all__ = ["HeteroGraphConv", "HeteroLinear", "HeteroEmbedding"]
class HeteroGraphConv(nn.Module):
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 torch as th
>>> h1 = {'user' : th.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 = torch.stack(tensors, dim=0)
return torch.sum(stacked, dim=0)
Attributes
----------
mods : dict[str, nn.Module]
Modules associated with every edge types.
"""
def __init__(self, mods, aggregate="sum"):
super(HeteroGraphConv, self).__init__()
self.mod_dict = mods
mods = {str(k): v for k, v in mods.items()}
# Register as child modules
self.mods = nn.ModuleDict(mods)
# PyTorch ModuleDict doesn't have get() method, so I have to store two
# dictionaries so that I can index with both canonical edge type and
# edge type with the get() method.
# 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 _get_module(self, etype):
mod = self.mod_dict.get(etype, None)
if mod is not None:
return mod
if isinstance(etype, tuple):
# etype is canonical
_, etype, _ = etype
return self.mod_dict[etype]
raise KeyError("Cannot find module with edge type %s" % etype)
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) or g.is_block:
if isinstance(inputs, tuple):
src_inputs, dst_inputs = inputs
else:
src_inputs = inputs
dst_inputs = {
k: v[: g.number_of_dst_nodes(k)] for k, v in inputs.items()
}
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._get_module((stype, etype, dtype))(
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._get_module((stype, etype, dtype))(
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 _max_reduce_func(inputs, dim):
return th.max(inputs, dim=dim)[0]
def _min_reduce_func(inputs, dim):
return th.min(inputs, dim=dim)[0]
def _sum_reduce_func(inputs, dim):
return th.sum(inputs, dim=dim)
def _mean_reduce_func(inputs, dim):
return th.mean(inputs, dim=dim)
def _stack_agg_func(inputs, dsttype): # pylint: disable=unused-argument
if len(inputs) == 0:
return None
return th.stack(inputs, dim=1)
def _agg_func(inputs, dsttype, fn): # pylint: disable=unused-argument
if len(inputs) == 0:
return None
stacked = th.stack(inputs, dim=0)
return fn(stacked, dim=0)
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 = _sum_reduce_func
elif agg == "max":
fn = _max_reduce_func
elif agg == "min":
fn = _min_reduce_func
elif agg == "mean":
fn = _mean_reduce_func
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":
return _stack_agg_func
else:
return partial(_agg_func, fn=fn)
class HeteroLinear(nn.Module):
"""Apply linear transformations on heterogeneous inputs.
Parameters
----------
in_size : dict[key, int]
Input feature size for heterogeneous inputs. A key can be a string or a tuple of strings.
out_size : int
Output feature size.
bias : bool, optional
If True, learns a bias term. Defaults: ``True``.
Examples
--------
>>> import dgl
>>> import torch
>>> from dgl.nn import HeteroLinear
>>> layer = HeteroLinear({'user': 1, ('user', 'follows', 'user'): 2}, 3)
>>> in_feats = {'user': torch.randn(2, 1), ('user', 'follows', 'user'): torch.randn(3, 2)}
>>> out_feats = layer(in_feats)
>>> print(out_feats['user'].shape)
torch.Size([2, 3])
>>> print(out_feats[('user', 'follows', 'user')].shape)
torch.Size([3, 3])
"""
def __init__(self, in_size, out_size, bias=True):
super(HeteroLinear, self).__init__()
self.linears = nn.ModuleDict()
for typ, typ_in_size in in_size.items():
self.linears[str(typ)] = nn.Linear(typ_in_size, out_size, bias=bias)
def forward(self, feat):
"""Forward function
Parameters
----------
feat : dict[key, Tensor]
Heterogeneous input features. It maps keys to features.
Returns
-------
dict[key, Tensor]
Transformed features.
"""
out_feat = dict()
for typ, typ_feat in feat.items():
out_feat[typ] = self.linears[str(typ)](typ_feat)
return out_feat
class HeteroEmbedding(nn.Module):
"""Create a heterogeneous embedding table.
It internally contains multiple ``torch.nn.Embedding`` with different dictionary sizes.
Parameters
----------
num_embeddings : dict[key, int]
Size of the dictionaries. A key can be a string or a tuple of strings.
embedding_dim : int
Size of each embedding vector.
Examples
--------
>>> import dgl
>>> import torch
>>> from dgl.nn import HeteroEmbedding
>>> layer = HeteroEmbedding({'user': 2, ('user', 'follows', 'user'): 3}, 4)
>>> # Get the heterogeneous embedding table
>>> embeds = layer.weight
>>> print(embeds['user'].shape)
torch.Size([2, 4])
>>> print(embeds[('user', 'follows', 'user')].shape)
torch.Size([3, 4])
>>> # Get the embeddings for a subset
>>> input_ids = {'user': torch.LongTensor([0]),
... ('user', 'follows', 'user'): torch.LongTensor([0, 2])}
>>> embeds = layer(input_ids)
>>> print(embeds['user'].shape)
torch.Size([1, 4])
>>> print(embeds[('user', 'follows', 'user')].shape)
torch.Size([2, 4])
"""
def __init__(self, num_embeddings, embedding_dim):
super(HeteroEmbedding, self).__init__()
self.embeds = nn.ModuleDict()
self.raw_keys = dict()
for typ, typ_num_rows in num_embeddings.items():
self.embeds[str(typ)] = nn.Embedding(typ_num_rows, embedding_dim)
self.raw_keys[str(typ)] = typ
@property
def weight(self):
"""Get the heterogeneous embedding table
Returns
-------
dict[key, Tensor]
Heterogeneous embedding table
"""
return {
self.raw_keys[typ]: emb.weight for typ, emb in self.embeds.items()
}
def reset_parameters(self):
"""
Use the xavier method in nn.init module to make the parameters uniformly distributed
"""
for typ in self.embeds.keys():
nn.init.xavier_uniform_(self.embeds[typ].weight)
def forward(self, input_ids):
"""Forward function
Parameters
----------
input_ids : dict[key, Tensor]
The row IDs to retrieve embeddings. It maps a key to key-specific IDs.
Returns
-------
dict[key, Tensor]
The retrieved embeddings.
"""
embeds = dict()
for typ, typ_ids in input_ids.items():
embeds[typ] = self.embeds[str(typ)](typ_ids)
return embeds
python/dgl/nn/pytorch/linear.py
+223
View File
@@ -0,0 +1,223 @@
"""Various commonly used linear modules"""
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
import math
import torch
import torch.nn as nn
from ...ops import gather_mm, segment_mm
__all__ = ["TypedLinear"]
class TypedLinear(nn.Module):
r"""Linear transformation according to types.
For each sample of the input batch :math:`x \in X`, apply linear transformation
:math:`xW_t`, where :math:`t` is the type of :math:`x`.
The module supports two regularization methods (basis-decomposition and
block-diagonal-decomposition) proposed by "`Modeling Relational Data
with Graph Convolutional Networks <https://arxiv.org/abs/1703.06103>`__"
The basis regularization decomposes :math:`W_t` by:
.. math::
W_t^{(l)} = \sum_{b=1}^B a_{tb}^{(l)}V_b^{(l)}
where :math:`B` is the number of bases, :math:`V_b^{(l)}` are linearly combined
with coefficients :math:`a_{tb}^{(l)}`.
The block-diagonal-decomposition regularization decomposes :math:`W_t` into :math:`B`
block-diagonal matrices. We refer to :math:`B` as the number of bases:
.. math::
W_t^{(l)} = \oplus_{b=1}^B Q_{tb}^{(l)}
where :math:`B` is the number of bases, :math:`Q_{tb}^{(l)}` are block
bases with shape :math:`R^{(d^{(l+1)}/B)\times(d^{l}/B)}`.
Parameters
----------
in_size : int
Input feature size.
out_size : int
Output feature size.
num_types : int
Total number of types.
regularizer : str, optional
Which weight regularizer to use "basis" or "bdd":
- "basis" is short for basis-decomposition.
- "bdd" is short for block-diagonal-decomposition.
Default applies no regularization.
num_bases : int, optional
Number of bases. Needed when ``regularizer`` is specified. Typically smaller
than ``num_types``.
Default: ``None``.
Examples
--------
No regularization.
>>> from dgl.nn import TypedLinear
>>> import torch
>>>
>>> x = torch.randn(100, 32)
>>> x_type = torch.randint(0, 5, (100,))
>>> m = TypedLinear(32, 64, 5)
>>> y = m(x, x_type)
>>> print(y.shape)
torch.Size([100, 64])
With basis regularization
>>> x = torch.randn(100, 32)
>>> x_type = torch.randint(0, 5, (100,))
>>> m = TypedLinear(32, 64, 5, regularizer='basis', num_bases=4)
>>> y = m(x, x_type)
>>> print(y.shape)
torch.Size([100, 64])
"""
def __init__(
self, in_size, out_size, num_types, regularizer=None, num_bases=None
):
super().__init__()
self.in_size = in_size
self.out_size = out_size
self.num_types = num_types
if regularizer is None:
self.W = nn.Parameter(torch.Tensor(num_types, in_size, out_size))
elif regularizer == "basis":
if num_bases is None:
raise ValueError(
'Missing "num_bases" for basis regularization.'
)
self.W = nn.Parameter(torch.Tensor(num_bases, in_size, out_size))
self.coeff = nn.Parameter(torch.Tensor(num_types, num_bases))
self.num_bases = num_bases
elif regularizer == "bdd":
if num_bases is None:
raise ValueError('Missing "num_bases" for bdd regularization.')
if in_size % num_bases != 0 or out_size % num_bases != 0:
raise ValueError(
"Input and output sizes must be divisible by num_bases."
)
self.submat_in = in_size // num_bases
self.submat_out = out_size // num_bases
self.W = nn.Parameter(
torch.Tensor(
num_types, num_bases * self.submat_in * self.submat_out
)
)
self.num_bases = num_bases
else:
raise ValueError(
f'Supported regularizer options: "basis", "bdd", but got {regularizer}'
)
self.regularizer = regularizer
self.reset_parameters()
def reset_parameters(self):
"""Reset parameters"""
with torch.no_grad():
# Follow torch.nn.Linear 's initialization to use kaiming_uniform_ on in_size
if self.regularizer is None:
nn.init.uniform_(
self.W,
-1 / math.sqrt(self.in_size),
1 / math.sqrt(self.in_size),
)
elif self.regularizer == "basis":
nn.init.uniform_(
self.W,
-1 / math.sqrt(self.in_size),
1 / math.sqrt(self.in_size),
)
nn.init.xavier_uniform_(
self.coeff, gain=nn.init.calculate_gain("relu")
)
elif self.regularizer == "bdd":
nn.init.uniform_(
self.W,
-1 / math.sqrt(self.submat_in),
1 / math.sqrt(self.submat_in),
)
else:
raise ValueError(
f'Supported regularizer options: "basis", "bdd", but got {regularizer}'
)
def get_weight(self):
"""Get type-wise weight"""
if self.regularizer is None:
return self.W
elif self.regularizer == "basis":
W = self.W.view(self.num_bases, self.in_size * self.out_size)
return (self.coeff @ W).view(
self.num_types, self.in_size, self.out_size
)
elif self.regularizer == "bdd":
return self.W
else:
raise ValueError(
f'Supported regularizer options: "basis", "bdd", but got {regularizer}'
)
def forward(self, x, x_type, sorted_by_type=False):
"""Forward computation.
Parameters
----------
x : torch.Tensor
A 2D input tensor. Shape: (N, D1)
x_type : torch.Tensor
A 1D integer tensor storing the type of the elements in ``x`` with one-to-one
correspondenc. Shape: (N,)
sorted_by_type : bool, optional
Whether the inputs have been sorted by the types. Forward on pre-sorted inputs may
be faster.
Returns
-------
y : torch.Tensor
The transformed output tensor. Shape: (N, D2)
"""
w = self.get_weight()
if self.regularizer == "bdd":
w = w.index_select(0, x_type).view(
-1, self.submat_in, self.submat_out
)
x = x.view(-1, 1, self.submat_in)
return torch.bmm(x, w).view(-1, self.out_size)
elif sorted_by_type:
pos_l = torch.searchsorted(
x_type, torch.arange(self.num_types, device=x.device)
)
pos_r = torch.cat(
[pos_l[1:], torch.tensor([len(x_type)], device=x.device)]
)
seglen = (
pos_r - pos_l
).cpu() # XXX(minjie): cause device synchronize
return segment_mm(x, w, seglen_a=seglen)
else:
return gather_mm(x, w, idx_b=x_type)
def __repr__(self):
if self.regularizer is None:
return (
f"TypedLinear(in_size={self.in_size}, out_size={self.out_size}, "
f"num_types={self.num_types})"
)
else:
return (
f"TypedLinear(in_size={self.in_size}, out_size={self.out_size}, "
f"num_types={self.num_types}, regularizer={self.regularizer}, "
f"num_bases={self.num_bases})"
)
python/dgl/nn/pytorch/link/__init__.py
+5
View File
@@ -0,0 +1,5 @@
"""Torch modules for link prediction/knowledge graph completion."""
from .edgepred import EdgePredictor
from .transe import TransE
from .transr import TransR
python/dgl/nn/pytorch/link/edgepred.py
+172
View File
@@ -0,0 +1,172 @@
"""Predictor for edges in homogeneous graphs."""
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
import torch
import torch.nn as nn
import torch.nn.functional as F
class EdgePredictor(nn.Module):
r"""Predictor/score function for pairs of node representations
Given a pair of node representations, :math:`h_i` and :math:`h_j`, it combines them with
**dot product**
.. math::
h_i^{T} h_j
or **cosine similarity**
.. math::
\frac{h_i^{T} h_j}{{\| h_i \|}_2 \cdot {\| h_j \|}_2}
or **elementwise product**
.. math::
h_i \odot h_j
or **concatenation**
.. math::
h_i \Vert h_j
Optionally, it passes the combined results to a linear layer for the final prediction.
Parameters
----------
op : str
The operation to apply. It can be 'dot', 'cos', 'ele', or 'cat',
corresponding to the equations above in order.
in_feats : int, optional
The input feature size of :math:`h_i` and :math:`h_j`. It is required
only if a linear layer is to be applied.
out_feats : int, optional
The output feature size. It is reuiqred only if a linear layer is to be applied.
bias : bool, optional
Whether to use bias for the linear layer if it applies.
Examples
--------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import EdgePredictor
>>> num_nodes = 2
>>> num_edges = 3
>>> in_feats = 4
>>> g = dgl.rand_graph(num_nodes=num_nodes, num_edges=num_edges)
>>> h = th.randn(num_nodes, in_feats)
>>> src, dst = g.edges()
>>> h_src = h[src]
>>> h_dst = h[dst]
Case1: dot product
>>> predictor = EdgePredictor('dot')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 1])
>>> predictor = EdgePredictor('dot', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])
Case2: cosine similarity
>>> predictor = EdgePredictor('cos')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 1])
>>> predictor = EdgePredictor('cos', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])
Case3: elementwise product
>>> predictor = EdgePredictor('ele')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 4])
>>> predictor = EdgePredictor('ele', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])
Case4: concatenation
>>> predictor = EdgePredictor('cat')
>>> predictor(h_src, h_dst).shape
torch.Size([3, 8])
>>> predictor = EdgePredictor('cat', in_feats, out_feats=3)
>>> predictor.reset_parameters()
>>> predictor(h_src, h_dst).shape
torch.Size([3, 3])
"""
def __init__(self, op, in_feats=None, out_feats=None, bias=False):
super(EdgePredictor, self).__init__()
assert op in [
"dot",
"cos",
"ele",
"cat",
], "Expect op to be in ['dot', 'cos', 'ele', 'cat'], got {}".format(op)
self.op = op
if (in_feats is not None) and (out_feats is not None):
if op in ["dot", "cos"]:
in_feats = 1
elif op == "cat":
in_feats = 2 * in_feats
self.linear = nn.Linear(in_feats, out_feats, bias=bias)
else:
self.linear = None
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
"""
if self.linear is not None:
self.linear.reset_parameters()
def forward(self, h_src, h_dst):
r"""
Description
-----------
Predict for pairs of node representations.
Parameters
----------
h_src : torch.Tensor
Source node features. The tensor is of shape :math:`(E, D_{in})`,
where :math:`E` is the number of edges/node pairs, and :math:`D_{in}`
is the input feature size.
h_dst : torch.Tensor
Destination node features. The tensor is of shape :math:`(E, D_{in})`,
where :math:`E` is the number of edges/node pairs, and :math:`D_{in}`
is the input feature size.
Returns
-------
torch.Tensor
The output features.
"""
if self.op == "dot":
N, D = h_src.shape
h = torch.bmm(h_src.view(N, 1, D), h_dst.view(N, D, 1)).squeeze(-1)
elif self.op == "cos":
h = F.cosine_similarity(h_src, h_dst).unsqueeze(-1)
elif self.op == "ele":
h = h_src * h_dst
else:
h = torch.cat([h_src, h_dst], dim=-1)
if self.linear is not None:
h = self.linear(h)
return h
python/dgl/nn/pytorch/link/transe.py
+99
View File
@@ -0,0 +1,99 @@
"""TransE."""
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
import torch
import torch.nn as nn
class TransE(nn.Module):
r"""Similarity measure from `Translating Embeddings for Modeling Multi-relational Data
<https://papers.nips.cc/paper/2013/hash/1cecc7a77928ca8133fa24680a88d2f9-Abstract.html>`__
Mathematically, it is defined as follows:
.. math::
- {\| h + r - t \|}_p
where :math:`h` is the head embedding, :math:`r` is the relation embedding, and
:math:`t` is the tail embedding.
Parameters
----------
num_rels : int
Number of relation types.
feats : int
Embedding size.
p : int, optional
The p to use for Lp norm, which can be 1 or 2.
Attributes
----------
rel_emb : torch.nn.Embedding
The learnable relation type embedding.
Examples
--------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import TransE
>>> # input features
>>> num_nodes = 10
>>> num_edges = 30
>>> num_rels = 3
>>> feats = 4
>>> scorer = TransE(num_rels=num_rels, feats=feats)
>>> g = dgl.rand_graph(num_nodes=num_nodes, num_edges=num_edges)
>>> src, dst = g.edges()
>>> h = th.randn(num_nodes, feats)
>>> h_head = h[src]
>>> h_tail = h[dst]
>>> # Randomly initialize edge relation types for demonstration
>>> rels = th.randint(low=0, high=num_rels, size=(num_edges,))
>>> scorer(h_head, h_tail, rels).shape
torch.Size([30])
"""
def __init__(self, num_rels, feats, p=1):
super(TransE, self).__init__()
self.rel_emb = nn.Embedding(num_rels, feats)
self.p = p
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
"""
self.rel_emb.reset_parameters()
def forward(self, h_head, h_tail, rels):
r"""
Description
-----------
Score triples.
Parameters
----------
h_head : torch.Tensor
Head entity features. The tensor is of shape :math:`(E, D)`, where
:math:`E` is the number of triples, and :math:`D` is the feature size.
h_tail : torch.Tensor
Tail entity features. The tensor is of shape :math:`(E, D)`, where
:math:`E` is the number of triples, and :math:`D` is the feature size.
rels : torch.Tensor
Relation types. It is a LongTensor of shape :math:`(E)`, where
:math:`E` is the number of triples.
Returns
-------
torch.Tensor
The triple scores. The tensor is of shape :math:`(E)`.
"""
h_rel = self.rel_emb(rels)
return -torch.norm(h_head + h_rel - h_tail, p=self.p, dim=-1)
python/dgl/nn/pytorch/link/transr.py
+108
View File
@@ -0,0 +1,108 @@
"""TransR."""
# pylint: disable= no-member, arguments-differ, invalid-name, W0235
import torch
import torch.nn as nn
class TransR(nn.Module):
r"""Similarity measure from
`Learning entity and relation embeddings for knowledge graph completion
<https://ojs.aaai.org/index.php/AAAI/article/view/9491>`__
Mathematically, it is defined as follows:
.. math::
- {\| M_r h + r - M_r t \|}_p
where :math:`M_r` is a relation-specific projection matrix, :math:`h` is the
head embedding, :math:`r` is the relation embedding, and :math:`t` is the tail embedding.
Parameters
----------
num_rels : int
Number of relation types.
rfeats : int
Relation embedding size.
nfeats : int
Entity embedding size.
p : int, optional
The p to use for Lp norm, which can be 1 or 2.
Attributes
----------
rel_emb : torch.nn.Embedding
The learnable relation type embedding.
rel_project : torch.nn.Embedding
The learnable relation-type-specific projection.
Examples
--------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import TransR
>>> # input features
>>> num_nodes = 10
>>> num_edges = 30
>>> num_rels = 3
>>> feats = 4
>>> scorer = TransR(num_rels=num_rels, rfeats=2, nfeats=feats)
>>> g = dgl.rand_graph(num_nodes=num_nodes, num_edges=num_edges)
>>> src, dst = g.edges()
>>> h = th.randn(num_nodes, feats)
>>> h_head = h[src]
>>> h_tail = h[dst]
>>> # Randomly initialize edge relation types for demonstration
>>> rels = th.randint(low=0, high=num_rels, size=(num_edges,))
>>> scorer(h_head, h_tail, rels).shape
torch.Size([30])
"""
def __init__(self, num_rels, rfeats, nfeats, p=1):
super(TransR, self).__init__()
self.rel_emb = nn.Embedding(num_rels, rfeats)
self.rel_project = nn.Embedding(num_rels, nfeats * rfeats)
self.rfeats = rfeats
self.nfeats = nfeats
self.p = p
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters.
"""
self.rel_emb.reset_parameters()
self.rel_project.reset_parameters()
def forward(self, h_head, h_tail, rels):
r"""
Score triples.
Parameters
----------
h_head : torch.Tensor
Head entity features. The tensor is of shape :math:`(E, D)`, where
:math:`E` is the number of triples, and :math:`D` is the feature size.
h_tail : torch.Tensor
Tail entity features. The tensor is of shape :math:`(E, D)`, where
:math:`E` is the number of triples, and :math:`D` is the feature size.
rels : torch.Tensor
Relation types. It is a LongTensor of shape :math:`(E)`, where
:math:`E` is the number of triples.
Returns
-------
torch.Tensor
The triple scores. The tensor is of shape :math:`(E)`.
"""
h_rel = self.rel_emb(rels)
proj_rel = self.rel_project(rels).reshape(-1, self.nfeats, self.rfeats)
h_head = (h_head.unsqueeze(1) @ proj_rel).squeeze(1)
h_tail = (h_tail.unsqueeze(1) @ proj_rel).squeeze(1)
return -torch.norm(h_head + h_rel - h_tail, p=self.p, dim=-1)
python/dgl/nn/pytorch/network_emb.py
+444
View File
@@ -0,0 +1,444 @@
"""Network Embedding NN Modules"""
# pylint: disable= invalid-name
import random
import torch
import torch.nn.functional as F
from torch import nn
from torch.nn import init
from tqdm.auto import trange
from ...base import NID
from ...convert import to_heterogeneous, to_homogeneous
from ...random import choice
from ...sampling import random_walk
__all__ = ["DeepWalk", "MetaPath2Vec"]
class DeepWalk(nn.Module):
"""DeepWalk module from `DeepWalk: Online Learning of Social Representations
<https://arxiv.org/abs/1403.6652>`__
For a graph, it learns the node representations from scratch by maximizing the similarity of
node pairs that are nearby (positive node pairs) and minimizing the similarity of other
random node pairs (negative node pairs).
Parameters
----------
g : DGLGraph
Graph for learning node embeddings
emb_dim : int, optional
Size of each embedding vector. Default: 128
walk_length : int, optional
Number of nodes in a random walk sequence. Default: 40
window_size : int, optional
In a random walk :attr:`w`, a node :attr:`w[j]` is considered close to a node
:attr:`w[i]` if :attr:`i - window_size <= j <= i + window_size`. Default: 5
neg_weight : float, optional
Weight of the loss term for negative samples in the total loss. Default: 1.0
negative_size : int, optional
Number of negative samples to use for each positive sample. Default: 5
fast_neg : bool, optional
If True, it samples negative node pairs within a batch of random walks. Default: True
sparse : bool, optional
If True, gradients with respect to the learnable weights will be sparse.
Default: True
Attributes
----------
node_embed : nn.Embedding
Embedding table of the nodes
Examples
--------
>>> import torch
>>> from dgl.data import CoraGraphDataset
>>> from dgl.nn import DeepWalk
>>> from torch.optim import SparseAdam
>>> from torch.utils.data import DataLoader
>>> from sklearn.linear_model import LogisticRegression
>>> dataset = CoraGraphDataset()
>>> g = dataset[0]
>>> model = DeepWalk(g)
>>> dataloader = DataLoader(torch.arange(g.num_nodes()), batch_size=128,
... shuffle=True, collate_fn=model.sample)
>>> optimizer = SparseAdam(model.parameters(), lr=0.01)
>>> num_epochs = 5
>>> for epoch in range(num_epochs):
... for batch_walk in dataloader:
... loss = model(batch_walk)
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> train_mask = g.ndata['train_mask']
>>> test_mask = g.ndata['test_mask']
>>> X = model.node_embed.weight.detach()
>>> y = g.ndata['label']
>>> clf = LogisticRegression().fit(X[train_mask].numpy(), y[train_mask].numpy())
>>> clf.score(X[test_mask].numpy(), y[test_mask].numpy())
"""
def __init__(
self,
g,
emb_dim=128,
walk_length=40,
window_size=5,
neg_weight=1,
negative_size=5,
fast_neg=True,
sparse=True,
):
super().__init__()
assert (
walk_length >= window_size + 1
), f"Expect walk_length >= window_size + 1, got {walk_length} and {window_size + 1}"
self.g = g
self.emb_dim = emb_dim
self.window_size = window_size
self.walk_length = walk_length
self.neg_weight = neg_weight
self.negative_size = negative_size
self.fast_neg = fast_neg
num_nodes = g.num_nodes()
# center node embedding
self.node_embed = nn.Embedding(num_nodes, emb_dim, sparse=sparse)
self.context_embed = nn.Embedding(num_nodes, emb_dim, sparse=sparse)
self.reset_parameters()
if not fast_neg:
neg_prob = g.out_degrees().pow(0.75)
# categorical distribution for true negative sampling
self.neg_prob = neg_prob / neg_prob.sum()
# Get list index pairs for positive samples.
# Given i, positive index pairs are (i - window_size, i), ... ,
# (i - 1, i), (i + 1, i), ..., (i + window_size, i)
idx_list_src = []
idx_list_dst = []
for i in range(walk_length):
for j in range(max(0, i - window_size), i):
idx_list_src.append(j)
idx_list_dst.append(i)
for j in range(i + 1, min(walk_length, i + 1 + window_size)):
idx_list_src.append(j)
idx_list_dst.append(i)
self.idx_list_src = torch.LongTensor(idx_list_src)
self.idx_list_dst = torch.LongTensor(idx_list_dst)
def reset_parameters(self):
"""Reinitialize learnable parameters"""
init_range = 1.0 / self.emb_dim
init.uniform_(self.node_embed.weight.data, -init_range, init_range)
init.constant_(self.context_embed.weight.data, 0)
def sample(self, indices):
"""Sample random walks
Parameters
----------
indices : torch.Tensor
Nodes from which we perform random walk
Returns
-------
torch.Tensor
Random walks in the form of node ID sequences. The Tensor
is of shape :attr:`(len(indices), walk_length)`.
"""
return random_walk(self.g, indices, length=self.walk_length - 1)[0]
def forward(self, batch_walk):
"""Compute the loss for the batch of random walks
Parameters
----------
batch_walk : torch.Tensor
Random walks in the form of node ID sequences. The Tensor
is of shape :attr:`(batch_size, walk_length)`.
Returns
-------
torch.Tensor
Loss value
"""
batch_size = len(batch_walk)
device = batch_walk.device
batch_node_embed = self.node_embed(batch_walk).view(-1, self.emb_dim)
batch_context_embed = self.context_embed(batch_walk).view(
-1, self.emb_dim
)
batch_idx_list_offset = torch.arange(batch_size) * self.walk_length
batch_idx_list_offset = batch_idx_list_offset.unsqueeze(1)
idx_list_src = batch_idx_list_offset + self.idx_list_src.unsqueeze(0)
idx_list_dst = batch_idx_list_offset + self.idx_list_dst.unsqueeze(0)
idx_list_src = idx_list_src.view(-1).to(device)
idx_list_dst = idx_list_dst.view(-1).to(device)
pos_src_emb = batch_node_embed[idx_list_src]
pos_dst_emb = batch_context_embed[idx_list_dst]
neg_idx_list_src = idx_list_dst.unsqueeze(1) + torch.zeros(
self.negative_size
).unsqueeze(0).to(device)
neg_idx_list_src = neg_idx_list_src.view(-1)
neg_src_emb = batch_node_embed[neg_idx_list_src.long()]
if self.fast_neg:
neg_idx_list_dst = list(range(batch_size * self.walk_length)) * (
self.negative_size * self.window_size * 2
)
random.shuffle(neg_idx_list_dst)
neg_idx_list_dst = neg_idx_list_dst[: len(neg_idx_list_src)]
neg_idx_list_dst = torch.LongTensor(neg_idx_list_dst).to(device)
neg_dst_emb = batch_context_embed[neg_idx_list_dst]
else:
neg_dst = choice(
self.g.num_nodes(), size=len(neg_src_emb), prob=self.neg_prob
)
neg_dst_emb = self.context_embed(neg_dst.to(device))
pos_score = torch.sum(torch.mul(pos_src_emb, pos_dst_emb), dim=1)
pos_score = torch.clamp(pos_score, max=6, min=-6)
pos_score = torch.mean(-F.logsigmoid(pos_score))
neg_score = torch.sum(torch.mul(neg_src_emb, neg_dst_emb), dim=1)
neg_score = torch.clamp(neg_score, max=6, min=-6)
neg_score = (
torch.mean(-F.logsigmoid(-neg_score))
* self.negative_size
* self.neg_weight
)
return torch.mean(pos_score + neg_score)
class MetaPath2Vec(nn.Module):
r"""metapath2vec module from `metapath2vec: Scalable Representation Learning for
Heterogeneous Networks <https://dl.acm.org/doi/pdf/10.1145/3097983.3098036>`__
To achieve efficient optimization, we leverage the negative sampling technique for the
training process. Repeatedly for each node in meta-path, we treat it as the center node
and sample nearby positive nodes within context size and draw negative samples among all
types of nodes from all meta-paths. Then we can use the center-context paired nodes and
context-negative paired nodes to update the network.
Parameters
----------
g : DGLGraph
Graph for learning node embeddings. Two different canonical edge types
:attr:`(utype, etype, vtype)` are not allowed to have same :attr:`etype`.
metapath : list[str]
A sequence of edge types in the form of a string. It defines a new edge type by composing
multiple edge types in order. Note that the start node type and the end one are commonly
the same.
window_size : int
In a random walk :attr:`w`, a node :attr:`w[j]` is considered close to a node
:attr:`w[i]` if :attr:`i - window_size <= j <= i + window_size`.
emb_dim : int, optional
Size of each embedding vector. Default: 128
negative_size : int, optional
Number of negative samples to use for each positive sample. Default: 5
sparse : bool, optional
If True, gradients with respect to the learnable weights will be sparse.
Default: True
Attributes
----------
node_embed : nn.Embedding
Embedding table of all nodes
local_to_global_nid : dict[str, list]
Mapping from type-specific node IDs to global node IDs
Examples
--------
>>> import torch
>>> import dgl
>>> from torch.optim import SparseAdam
>>> from torch.utils.data import DataLoader
>>> from dgl.nn.pytorch import MetaPath2Vec
>>> # Define a model
>>> g = dgl.heterograph({
... ('user', 'uc', 'company'): dgl.rand_graph(100, 1000).edges(),
... ('company', 'cp', 'product'): dgl.rand_graph(100, 1000).edges(),
... ('company', 'cu', 'user'): dgl.rand_graph(100, 1000).edges(),
... ('product', 'pc', 'company'): dgl.rand_graph(100, 1000).edges()
... })
>>> model = MetaPath2Vec(g, ['uc', 'cu'], window_size=1)
>>> # Use the source node type of etype 'uc'
>>> dataloader = DataLoader(torch.arange(g.num_nodes('user')), batch_size=128,
... shuffle=True, collate_fn=model.sample)
>>> optimizer = SparseAdam(model.parameters(), lr=0.025)
>>> for (pos_u, pos_v, neg_v) in dataloader:
... loss = model(pos_u, pos_v, neg_v)
... optimizer.zero_grad()
... loss.backward()
... optimizer.step()
>>> # Get the embeddings of all user nodes
>>> user_nids = torch.LongTensor(model.local_to_global_nid['user'])
>>> user_emb = model.node_embed(user_nids)
"""
def __init__(
self,
g,
metapath,
window_size,
emb_dim=128,
negative_size=5,
sparse=True,
):
super().__init__()
assert (
len(metapath) + 1 >= window_size
), f"Expect len(metapath) >= window_size - 1, got {metapath} and {window_size}"
self.hg = g
self.emb_dim = emb_dim
self.metapath = metapath
self.window_size = window_size
self.negative_size = negative_size
# convert edge metapath to node metapath
# get initial source node type
src_type, _, _ = g.to_canonical_etype(metapath[0])
node_metapath = [src_type]
for etype in metapath:
_, _, dst_type = g.to_canonical_etype(etype)
node_metapath.append(dst_type)
self.node_metapath = node_metapath
# Convert the graph into a homogeneous one for global to local node ID mapping
g = to_homogeneous(g)
# Convert it back to the hetero one for local to global node ID mapping
hg = to_heterogeneous(g, self.hg.ntypes, self.hg.etypes)
local_to_global_nid = hg.ndata[NID]
for key, val in local_to_global_nid.items():
local_to_global_nid[key] = list(val.cpu().numpy())
self.local_to_global_nid = local_to_global_nid
num_nodes_total = hg.num_nodes()
node_frequency = torch.zeros(num_nodes_total)
# random walk
for idx in trange(hg.num_nodes(node_metapath[0])):
traces, _ = random_walk(g=hg, nodes=[idx], metapath=metapath)
for tr in traces.cpu().numpy():
tr_nids = [
self.local_to_global_nid[node_metapath[i]][tr[i]]
for i in range(len(tr))
]
node_frequency[torch.LongTensor(tr_nids)] += 1
neg_prob = node_frequency.pow(0.75)
self.neg_prob = neg_prob / neg_prob.sum()
# center node embedding
self.node_embed = nn.Embedding(
num_nodes_total, self.emb_dim, sparse=sparse
)
self.context_embed = nn.Embedding(
num_nodes_total, self.emb_dim, sparse=sparse
)
self.reset_parameters()
def reset_parameters(self):
"""Reinitialize learnable parameters"""
init_range = 1.0 / self.emb_dim
init.uniform_(self.node_embed.weight.data, -init_range, init_range)
init.constant_(self.context_embed.weight.data, 0)
def sample(self, indices):
"""Sample positive and negative samples
Parameters
----------
indices : torch.Tensor
Node IDs of the source node type from which we perform random walks
Returns
-------
torch.Tensor
Positive center nodes
torch.Tensor
Positive context nodes
torch.Tensor
Negative context nodes
"""
traces, _ = random_walk(
g=self.hg, nodes=indices, metapath=self.metapath
)
u_list = []
v_list = []
for tr in traces.cpu().numpy():
tr_nids = [
self.local_to_global_nid[self.node_metapath[i]][tr[i]]
for i in range(len(tr))
]
for i, u in enumerate(tr_nids):
for j, v in enumerate(
tr_nids[max(i - self.window_size, 0) : i + self.window_size]
):
if i == j:
continue
u_list.append(u)
v_list.append(v)
neg_v = choice(
self.hg.num_nodes(),
size=len(u_list) * self.negative_size,
prob=self.neg_prob,
).reshape(len(u_list), self.negative_size)
return torch.LongTensor(u_list), torch.LongTensor(v_list), neg_v
def forward(self, pos_u, pos_v, neg_v):
r"""Compute the loss for the batch of positive and negative samples
Parameters
----------
pos_u : torch.Tensor
Positive center nodes
pos_v : torch.Tensor
Positive context nodes
neg_v : torch.Tensor
Negative context nodes
Returns
-------
torch.Tensor
Loss value
"""
emb_u = self.node_embed(pos_u)
emb_v = self.context_embed(pos_v)
emb_neg_v = self.context_embed(neg_v)
score = torch.sum(torch.mul(emb_u, emb_v), dim=1)
score = torch.clamp(score, max=10, min=-10)
score = -F.logsigmoid(score)
neg_score = torch.bmm(emb_neg_v, emb_u.unsqueeze(2)).squeeze()
neg_score = torch.clamp(neg_score, max=10, min=-10)
neg_score = -torch.sum(F.logsigmoid(-neg_score), dim=1)
return torch.mean(score + neg_score)
python/dgl/nn/pytorch/softmax.py
+3
View File
@@ -0,0 +1,3 @@
"""Torch modules for graph related softmax."""
# pylint: disable= unused-import
from ..functional import edge_softmax
python/dgl/nn/pytorch/sparse_emb.py
+479
View File
@@ -0,0 +1,479 @@
"""Torch NodeEmbedding."""
from datetime import timedelta
import torch as th
from ...backend import pytorch as F
from ...cuda import nccl
from ...partition import NDArrayPartition
from ...utils import create_shared_mem_array, get_shared_mem_array
_STORE = None
class NodeEmbedding: # NodeEmbedding
"""Class for storing node embeddings.
The class is optimized for training large-scale node embeddings. It updates the embedding in
a sparse way and can scale to graphs with millions of nodes. It also supports partitioning
to multiple GPUs (on a single machine) for more acceleration. It does not support partitioning
across machines.
Currently, DGL provides two optimizers that work with this NodeEmbedding
class: ``SparseAdagrad`` and ``SparseAdam``.
The implementation is based on torch.distributed package. It depends on the pytorch
default distributed process group to collect multi-process information and uses
``torch.distributed.TCPStore`` to share meta-data information across multiple gpu processes.
It use the local address of '127.0.0.1:12346' to initialize the TCPStore.
NOTE: The support of NodeEmbedding is experimental.
Parameters
----------
num_embeddings : int
The number of embeddings. Currently, the number of embeddings has to be the same as
the number of nodes.
embedding_dim : int
The dimension size of embeddings.
name : str
The name of the embeddings. The name should uniquely identify the embeddings in the system.
init_func : callable, optional
The function to create the initial data. If the init function is not provided,
the values of the embeddings are initialized to zero.
device : th.device
Device to store the embeddings on.
parittion : NDArrayPartition
The partition to use to distributed the embeddings between
processes.
Examples
--------
Before launching multiple gpu processes
>>> def initializer(emb):
th.nn.init.xavier_uniform_(emb)
return emb
In each training process
>>> emb = dgl.nn.NodeEmbedding(g.num_nodes(), 10, 'emb', init_func=initializer)
>>> optimizer = dgl.optim.SparseAdam([emb], lr=0.001)
>>> for blocks in dataloader:
... ...
... feats = emb(nids, gpu_0)
... loss = F.sum(feats + 1, 0)
... loss.backward()
... optimizer.step()
"""
def __init__(
self,
num_embeddings,
embedding_dim,
name,
init_func=None,
device=None,
partition=None,
):
global _STORE
if device is None:
device = th.device("cpu")
# Check whether it is multi-gpu training or not.
if th.distributed.is_initialized():
rank = th.distributed.get_rank()
world_size = th.distributed.get_world_size()
else:
rank = -1
world_size = 0
self._rank = rank
self._world_size = world_size
self._store = None
self._comm = None
self._partition = partition
host_name = "127.0.0.1"
port = 12346
if rank >= 0:
# for multi-gpu training, setup a TCPStore for
# embeding status synchronization across GPU processes
if _STORE is None:
_STORE = th.distributed.TCPStore(
host_name,
port,
world_size,
rank == 0,
timedelta(seconds=10 * 60),
)
self._store = _STORE
# embeddings is stored in CPU memory.
if th.device(device) == th.device("cpu"):
if rank <= 0:
emb = create_shared_mem_array(
name, (num_embeddings, embedding_dim), th.float32
)
if init_func is not None:
emb = init_func(emb)
if rank == 0: # the master gpu process
for _ in range(1, world_size):
# send embs
self._store.set(name, name)
elif rank > 0:
# receive
self._store.wait([name])
emb = get_shared_mem_array(
name, (num_embeddings, embedding_dim), th.float32
)
self._tensor = emb
else: # embeddings is stored in GPU memory.
self._comm = True
if not self._partition:
# for communication we need a partition
self._partition = NDArrayPartition(
num_embeddings,
self._world_size if self._world_size > 0 else 1,
mode="remainder",
)
# create local tensors for the weights
local_size = self._partition.local_size(max(self._rank, 0))
# TODO(dlasalle): support 16-bit/half embeddings
emb = th.empty(
[local_size, embedding_dim],
dtype=th.float32,
requires_grad=False,
device=device,
)
if init_func:
emb = init_func(emb)
self._tensor = emb
self._num_embeddings = num_embeddings
self._embedding_dim = embedding_dim
self._name = name
self._optm_state = None # track optimizer state
self._trace = [] # track minibatch
def __call__(self, node_ids, device=th.device("cpu")):
"""
node_ids : th.tensor
Index of the embeddings to collect.
device : th.device
Target device to put the collected embeddings.
"""
if not self._comm:
# For embeddings stored on the CPU.
emb = self._tensor[node_ids].to(device)
else:
# For embeddings stored on the GPU.
# The following method is designed to perform communication
# across multiple GPUs and can handle situations where only one GPU
# is present gracefully, a.k.a. self._world_size == 1 or
# 0 (when th.distributed.is_initialized() is false).
emb = nccl.sparse_all_to_all_pull(
node_ids, self._tensor, self._partition
)
emb = emb.to(device)
if F.is_recording():
emb = F.attach_grad(emb)
self._trace.append((node_ids.to(device), emb))
return emb
@property
def store(self):
"""Return torch.distributed.TCPStore for
meta data sharing across processes.
Returns
-------
torch.distributed.TCPStore
KVStore used for meta data sharing.
"""
return self._store
@property
def partition(self):
"""Return the partition identifying how the tensor is split across
processes.
Returns
-------
String
The mode.
"""
return self._partition
@property
def rank(self):
"""Return rank of current process.
Returns
-------
int
The rank of current process.
"""
return self._rank
@property
def world_size(self):
"""Return world size of the pytorch distributed training env.
Returns
-------
int
The world size of the pytorch distributed training env.
"""
return self._world_size
@property
def name(self):
"""Return the name of NodeEmbedding.
Returns
-------
str
The name of NodeEmbedding.
"""
return self._name
@property
def num_embeddings(self):
"""Return the number of embeddings.
Returns
-------
int
The number of embeddings.
"""
return self._num_embeddings
@property
def embedding_dim(self):
"""Return the dimension of embeddings.
Returns
-------
int
The dimension of embeddings.
"""
return self._embedding_dim
def set_optm_state(self, state):
"""Store the optimizer related state tensor.
Parameters
----------
state : tuple of torch.Tensor
Optimizer related state.
"""
self._optm_state = state
@property
def optm_state(self):
"""Return the optimizer related state tensor.
Returns
-------
tuple of torch.Tensor
The optimizer related state.
"""
return self._optm_state
@property
def trace(self):
"""Return a trace of the indices of embeddings
used in the training step(s).
Returns
-------
[torch.Tensor]
The indices of embeddings used in the training step(s).
"""
return self._trace
def reset_trace(self):
"""Clean up the trace of the indices of embeddings
used in the training step(s).
"""
self._trace = []
@property
def weight(self):
"""Return the tensor storing the node embeddings
Returns
-------
torch.Tensor
The tensor storing the node embeddings
"""
return self._tensor
def all_set_embedding(self, values):
"""Set the values of the embedding. This method must be called by all
processes sharing the embedding with identical tensors for
:attr:`values`.
NOTE: This method must be called by all processes sharing the
embedding, or it may result in a deadlock.
Parameters
----------
values : Tensor
The global tensor to pull values from.
"""
if self._partition:
idxs = F.copy_to(
self._partition.get_local_indices(
max(self._rank, 0),
ctx=F.context(self._tensor),
),
F.context(values),
)
self._tensor[:] = F.copy_to(
F.gather_row(values, idxs), ctx=F.context(self._tensor)
)[:]
else:
if self._rank == 0:
self._tensor[:] = F.copy_to(
values, ctx=F.context(self._tensor)
)[:]
if th.distributed.is_initialized():
th.distributed.barrier()
def _all_get_tensor(self, shared_name, tensor, shape):
"""A helper function to get model-parallel tensors.
This method must and only need to be called in multi-GPU DDP training.
For now, it's only used in ``all_get_embedding`` and
``_all_get_optm_state``.
"""
# create a shared memory tensor
if self._rank == 0:
# root process creates shared memory
val = create_shared_mem_array(
shared_name,
shape,
tensor.dtype,
)
self._store.set(shared_name, shared_name)
else:
self._store.wait([shared_name])
val = get_shared_mem_array(
shared_name,
shape,
tensor.dtype,
)
# need to map indices and slice into existing tensor
idxs = self._partition.map_to_global(
F.arange(0, tensor.shape[0], ctx=F.context(tensor)),
self._rank,
).to(val.device)
val[idxs] = tensor.to(val.device)
self._store.delete_key(shared_name)
# wait for all processes to finish
th.distributed.barrier()
return val
def all_get_embedding(self):
"""Return a copy of the embedding stored in CPU memory. If this is a
multi-processing instance, the tensor will be returned in shared
memory. If the embedding is currently stored on multiple GPUs, all
processes must call this method in the same order.
NOTE: This method must be called by all processes sharing the
embedding, or it may result in a deadlock.
Returns
-------
torch.Tensor
The tensor storing the node embeddings.
"""
if self._partition:
if self._world_size == 0:
# non-multiprocessing
return self._tensor.to(th.device("cpu"))
else:
return self._all_get_tensor(
f"{self._name}_gather",
self._tensor,
(self._num_embeddings, self._embedding_dim),
)
else:
# already stored in CPU memory
return self._tensor
def _all_get_optm_state(self):
"""Return a copy of the whole optimizer states stored in CPU memory.
If this is a multi-processing instance, the states will be returned in
shared memory. If the embedding is currently stored on multiple GPUs,
all processes must call this method in the same order.
NOTE: This method must be called by all processes sharing the
embedding, or it may result in a deadlock.
Returns
-------
tuple of torch.Tensor
The optimizer states stored in CPU memory.
"""
if self._partition:
if self._world_size == 0:
# non-multiprocessing
return tuple(
state.to(th.device("cpu")) for state in self._optm_state
)
else:
return tuple(
self._all_get_tensor(
f"state_gather_{self._name}_{i}",
state,
(self._num_embeddings, *state.shape[1:]),
)
for i, state in enumerate(self._optm_state)
)
else:
# already stored in CPU memory
return self._optm_state
def _all_set_optm_state(self, states):
"""Set the optimizer states of the embedding. This method must be
called by all processes sharing the embedding with identical
:attr:`states`.
NOTE: This method must be called by all processes sharing the
embedding, or it may result in a deadlock.
Parameters
----------
states : tuple of torch.Tensor
The global states to pull values from.
"""
if self._partition:
idxs = F.copy_to(
self._partition.get_local_indices(
max(self._rank, 0), ctx=F.context(self._tensor)
),
F.context(states[0]),
)
for state, new_state in zip(self._optm_state, states):
state[:] = F.copy_to(
F.gather_row(new_state, idxs), ctx=F.context(self._tensor)
)[:]
else:
# stored in CPU memory
if self._rank <= 0:
for state, new_state in zip(self._optm_state, states):
state[:] = F.copy_to(
new_state, ctx=F.context(self._tensor)
)[:]
if th.distributed.is_initialized():
th.distributed.barrier()
python/dgl/nn/pytorch/utils.py
+556
View File
@@ -0,0 +1,556 @@
"""Utilities for pytorch NN package"""
# pylint: disable=no-member, invalid-name
import torch as th
import torch.nn.functional as F
from torch import nn
from ... import DGLGraph, function as fn
from ...base import dgl_warning
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 : torch.Tensor
lhs tensor
B : torch.Tensor
rhs tensor
Returns
-------
C : torch.Tensor
result tensor
"""
if A.dtype == th.int64 and len(A.shape) == 1:
return B.index_select(0, A)
else:
return th.matmul(A, B)
def bmm_maybe_select(A, B, index):
"""Slice submatrices of A by the given index and perform bmm.
B is a 3D tensor of shape (N, D1, D2), which can be viewed as a stack of
N matrices of shape (D1, D2). The input index is an integer vector of length M.
A could be either:
(1) a dense tensor of shape (M, D1),
(2) an integer vector of length M.
The result C is a 2D matrix of shape (M, D2)
For case (1), C is computed by bmm:
::
C[i, :] = matmul(A[i, :], B[index[i], :, :])
For case (2), C is computed by index select:
::
C[i, :] = B[index[i], A[i], :]
Parameters
----------
A : torch.Tensor
lhs tensor
B : torch.Tensor
rhs tensor
index : torch.Tensor
index tensor
Returns
-------
C : torch.Tensor
return tensor
"""
if A.dtype == th.int64 and len(A.shape) == 1:
# following is a faster version of B[index, A, :]
B = B.view(-1, B.shape[2])
flatidx = index * B.shape[1] + A
return B.index_select(0, flatidx)
else:
BB = B.index_select(0, index)
return th.bmm(A.unsqueeze(1), BB).squeeze()
# pylint: disable=W0235
class Identity(nn.Module):
"""A placeholder identity operator that is argument-insensitive.
(Identity has already been supported by PyTorch 1.2, we will directly
import torch.nn.Identity in the future)
"""
def __init__(self):
super(Identity, self).__init__()
def forward(self, x):
"""Return input"""
return x
class Sequential(nn.Sequential):
r"""A sequential container for stacking graph neural network modules
DGL supports two modes: sequentially apply GNN modules on 1) the same graph or
2) a list of given graphs. In the second case, the number of graphs equals the
number of modules inside this container.
Parameters
----------
*args :
Sub-modules of torch.nn.Module that will be added to the container in
the order by which they are passed in the constructor.
Examples
--------
The following example uses PyTorch backend.
Mode 1: sequentially apply GNN modules on the same graph
>>> import torch
>>> import dgl
>>> import torch.nn as nn
>>> import dgl.function as fn
>>> from dgl.nn.pytorch import Sequential
>>> class ExampleLayer(nn.Module):
>>> def __init__(self):
>>> super().__init__()
>>> 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(ExampleLayer(), ExampleLayer(), ExampleLayer())
>>> n_feat = torch.rand(3, 4)
>>> e_feat = torch.rand(9, 4)
>>> net(g, n_feat, e_feat)
(tensor([[39.8597, 45.4542, 25.1877, 30.8086],
[40.7095, 45.3985, 25.4590, 30.0134],
[40.7894, 45.2556, 25.5221, 30.4220]]),
tensor([[80.3772, 89.7752, 50.7762, 60.5520],
[80.5671, 89.3736, 50.6558, 60.6418],
[80.4620, 89.5142, 50.3643, 60.3126],
[80.4817, 89.8549, 50.9430, 59.9108],
[80.2284, 89.6954, 50.0448, 60.1139],
[79.7846, 89.6882, 50.5097, 60.6213],
[80.2654, 90.2330, 50.2787, 60.6937],
[80.3468, 90.0341, 50.2062, 60.2659],
[80.0556, 90.2789, 50.2882, 60.5845]]))
Mode 2: sequentially apply GNN modules on different graphs
>>> import torch
>>> import dgl
>>> import torch.nn as nn
>>> import dgl.function as fn
>>> import networkx as nx
>>> from dgl.nn.pytorch import Sequential
>>> class ExampleLayer(nn.Module):
>>> def __init__(self):
>>> super().__init__()
>>> 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.view(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(ExampleLayer(), ExampleLayer(), ExampleLayer())
>>> n_feat = torch.rand(32, 4)
>>> net([g1, g2, g3], n_feat)
tensor([[209.6221, 225.5312, 193.8920, 220.1002],
[250.0169, 271.9156, 240.2467, 267.7766],
[220.4007, 239.7365, 213.8648, 234.9637],
[196.4630, 207.6319, 184.2927, 208.7465]])
"""
def __init__(self, *args):
super(Sequential, self).__init__(*args)
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 the :math:`i`-th module should match the input
of the :math:`(i+1)`-th module in the sequential.
"""
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
class WeightBasis(nn.Module):
r"""Basis decomposition from `Modeling Relational Data with Graph
Convolutional Networks <https://arxiv.org/abs/1703.06103>`__
It can be described as below:
.. math::
W_o = \sum_{b=1}^B a_{ob} V_b
Each weight output :math:`W_o` is essentially a linear combination of basis
transformations :math:`V_b` with coefficients :math:`a_{ob}`.
If is useful as a form of regularization on a large parameter matrix. Thus,
the number of weight outputs is usually larger than the number of bases.
Parameters
----------
shape : tuple[int]
Shape of the basis parameter.
num_bases : int
Number of bases.
num_outputs : int
Number of outputs.
"""
def __init__(self, shape, num_bases, num_outputs):
super(WeightBasis, self).__init__()
self.shape = shape
self.num_bases = num_bases
self.num_outputs = num_outputs
if num_outputs <= num_bases:
dgl_warning(
"The number of weight outputs should be larger than the number"
" of bases."
)
self.weight = nn.Parameter(th.Tensor(self.num_bases, *shape))
nn.init.xavier_uniform_(
self.weight, gain=nn.init.calculate_gain("relu")
)
# linear combination coefficients
self.w_comp = nn.Parameter(th.Tensor(self.num_outputs, self.num_bases))
nn.init.xavier_uniform_(
self.w_comp, gain=nn.init.calculate_gain("relu")
)
def forward(self):
r"""Forward computation
Returns
-------
weight : torch.Tensor
Composed weight tensor of shape ``(num_outputs,) + shape``
"""
# generate all weights from bases
weight = th.matmul(self.w_comp, self.weight.view(self.num_bases, -1))
return weight.view(self.num_outputs, *self.shape)
class JumpingKnowledge(nn.Module):
r"""The Jumping Knowledge aggregation module from `Representation Learning on
Graphs with Jumping Knowledge Networks <https://arxiv.org/abs/1806.03536>`__
It aggregates the output representations of multiple GNN layers with
**concatenation**
.. math::
h_i^{(1)} \, \Vert \, \ldots \, \Vert \, h_i^{(T)}
or **max pooling**
.. math::
\max \left( h_i^{(1)}, \ldots, h_i^{(T)} \right)
or **LSTM**
.. math::
\sum_{t=1}^T \alpha_i^{(t)} h_i^{(t)}
with attention scores :math:`\alpha_i^{(t)}` obtained from a BiLSTM
Parameters
----------
mode : str
The aggregation to apply. It can be 'cat', 'max', or 'lstm',
corresponding to the equations above in order.
in_feats : int, optional
This argument is only required if :attr:`mode` is ``'lstm'``.
The output representation size of a single GNN layer. Note that
all GNN layers need to have the same output representation size.
num_layers : int, optional
This argument is only required if :attr:`mode` is ``'lstm'``.
The number of GNN layers for output aggregation.
Examples
--------
>>> import dgl
>>> import torch as th
>>> from dgl.nn import JumpingKnowledge
>>> # Output representations of two GNN layers
>>> num_nodes = 3
>>> in_feats = 4
>>> feat_list = [th.zeros(num_nodes, in_feats), th.ones(num_nodes, in_feats)]
>>> # Case1
>>> model = JumpingKnowledge()
>>> model(feat_list).shape
torch.Size([3, 8])
>>> # Case2
>>> model = JumpingKnowledge(mode='max')
>>> model(feat_list).shape
torch.Size([3, 4])
>>> # Case3
>>> model = JumpingKnowledge(mode='max', in_feats=in_feats, num_layers=len(feat_list))
>>> model(feat_list).shape
torch.Size([3, 4])
"""
def __init__(self, mode="cat", in_feats=None, num_layers=None):
super(JumpingKnowledge, self).__init__()
assert mode in [
"cat",
"max",
"lstm",
], "Expect mode to be 'cat', or 'max' or 'lstm', got {}".format(mode)
self.mode = mode
if mode == "lstm":
assert in_feats is not None, "in_feats is required for lstm mode"
assert (
num_layers is not None
), "num_layers is required for lstm mode"
hidden_size = (num_layers * in_feats) // 2
self.lstm = nn.LSTM(
in_feats, hidden_size, bidirectional=True, batch_first=True
)
self.att = nn.Linear(2 * hidden_size, 1)
def reset_parameters(self):
r"""
Description
-----------
Reinitialize learnable parameters. This comes into effect only for the lstm mode.
"""
if self.mode == "lstm":
self.lstm.reset_parameters()
self.att.reset_parameters()
def forward(self, feat_list):
r"""
Description
-----------
Aggregate output representations across multiple GNN layers.
Parameters
----------
feat_list : list[Tensor]
feat_list[i] is the output representations of a GNN layer.
Returns
-------
Tensor
The aggregated representations.
"""
if self.mode == "cat":
return th.cat(feat_list, dim=-1)
elif self.mode == "max":
return th.stack(feat_list, dim=-1).max(dim=-1)[0]
else:
# LSTM
stacked_feat_list = th.stack(
feat_list, dim=1
) # (N, num_layers, in_feats)
alpha, _ = self.lstm(stacked_feat_list)
alpha = self.att(alpha).squeeze(-1) # (N, num_layers)
alpha = th.softmax(alpha, dim=-1)
return (stacked_feat_list * alpha.unsqueeze(-1)).sum(dim=1)
class LabelPropagation(nn.Module):
r"""Label Propagation from `Learning from Labeled and Unlabeled Data with Label
Propagation <http://mlg.eng.cam.ac.uk/zoubin/papers/CMU-CALD-02-107.pdf>`__
.. math::
\mathbf{Y}^{(t+1)} = \alpha \tilde{A} \mathbf{Y}^{(t)} + (1 - \alpha) \mathbf{Y}^{(0)}
where unlabeled data is initially set to zero and inferred from labeled data via
propagation. :math:`\alpha` is a weight parameter for balancing between updated labels
and initial labels. :math:`\tilde{A}` denotes the normalized adjacency matrix.
Parameters
----------
k: int
The number of propagation steps.
alpha : float
The :math:`\alpha` coefficient in range [0, 1].
norm_type : str, optional
The type of normalization applied to the adjacency matrix, must be one of the
following choices:
* ``row``: row-normalized adjacency as :math:`D^{-1}A`
* ``sym``: symmetrically normalized adjacency as :math:`D^{-1/2}AD^{-1/2}`
Default: 'sym'.
clamp : bool, optional
A bool flag to indicate whether to clamp the labels to [0, 1] after propagation.
Default: True.
normalize: bool, optional
A bool flag to indicate whether to apply row-normalization after propagation.
Default: False.
reset : bool, optional
A bool flag to indicate whether to reset the known labels after each
propagation step. Default: False.
Examples
--------
>>> import torch
>>> import dgl
>>> from dgl.nn import LabelPropagation
>>> label_propagation = LabelPropagation(k=5, alpha=0.5, clamp=False, normalize=True)
>>> g = dgl.rand_graph(5, 10)
>>> labels = torch.tensor([0, 2, 1, 3, 0]).long()
>>> mask = torch.tensor([0, 1, 1, 1, 0]).bool()
>>> new_labels = label_propagation(g, labels, mask)
"""
def __init__(
self,
k,
alpha,
norm_type="sym",
clamp=True,
normalize=False,
reset=False,
):
super(LabelPropagation, self).__init__()
self.k = k
self.alpha = alpha
self.norm_type = norm_type
self.clamp = clamp
self.normalize = normalize
self.reset = reset
def forward(self, g, labels, mask=None):
r"""Compute the label propagation process.
Parameters
----------
g : DGLGraph
The input graph.
labels : torch.Tensor
The input node labels. There are three cases supported.
* A LongTensor of shape :math:`(N, 1)` or :math:`(N,)` for node class labels in
multiclass classification, where :math:`N` is the number of nodes.
* A LongTensor of shape :math:`(N, C)` for one-hot encoding of node class labels
in multiclass classification, where :math:`C` is the number of classes.
* A LongTensor of shape :math:`(N, L)` for node labels in multilabel binary
classification, where :math:`L` is the number of labels.
mask : torch.Tensor
The bool indicators of shape :math:`(N,)` with True denoting labeled nodes.
Default: None, indicating all nodes are labeled.
Returns
-------
torch.Tensor
The propagated node labels of shape :math:`(N, D)` with float type, where :math:`D`
is the number of classes or labels.
"""
with g.local_scope():
# multi-label / multi-class
if len(labels.size()) > 1 and labels.size(1) > 1:
labels = labels.to(th.float32)
# single-label multi-class
else:
labels = F.one_hot(labels.view(-1)).to(th.float32)
y = labels
if mask is not None:
y = th.zeros_like(labels)
y[mask] = labels[mask]
init = (1 - self.alpha) * y
in_degs = g.in_degrees().float().clamp(min=1)
out_degs = g.out_degrees().float().clamp(min=1)
if self.norm_type == "sym":
norm_i = th.pow(in_degs, -0.5).to(labels.device).unsqueeze(1)
norm_j = th.pow(out_degs, -0.5).to(labels.device).unsqueeze(1)
elif self.norm_type == "row":
norm_i = th.pow(in_degs, -1.0).to(labels.device).unsqueeze(1)
else:
raise ValueError(
f"Expect norm_type to be 'sym' or 'row', got {self.norm_type}"
)
for _ in range(self.k):
g.ndata["h"] = y * norm_j if self.norm_type == "sym" else y
g.update_all(fn.copy_u("h", "m"), fn.sum("m", "h"))
y = init + self.alpha * g.ndata["h"] * norm_i
if self.clamp:
y = y.clamp_(0.0, 1.0)
if self.normalize:
y = F.normalize(y, p=1)
if self.reset:
y[mask] = labels[mask]
return y