chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user