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 and both. 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