wehub-resource-sync
700 lines
19 KiB
Python
700 lines
19 KiB
Python
"""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
|