DegreeEncoder¶
-
class
dgl.nn.pytorch.graph_transformer.
DegreeEncoder
(max_degree, embedding_dim, direction='both')[source]¶ Bases:
torch.nn.modules.module.Module
Degree Encoder, as introduced in Do Transformers Really Perform Bad for Graph Representation? 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
andboth
.both
encodes degrees from both directions and output the addition of them. Default :both
.
Example
>>> import dgl >>> from dgl.nn import DegreeEncoder
>>> g = dgl.graph(([0,0,0,1,1,2,3,3], [1,2,3,0,3,0,0,1])) >>> degree_encoder = DegreeEncoder(5, 16) >>> degree_embedding = degree_encoder(g)
-
forward
(g)[source]¶ - Parameters
g (DGLGraph) – A DGLGraph to be encoded. If it is a heterogeneous one, it will be transformed into a homogeneous one first.
- Returns
Return degree embedding vectors of shape \((N, embedding_dim)\), where \(N\) is th number of nodes in the input graph.
- Return type
Tensor