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