DegreeEncoder

class dgl.nn.pytorch.gt.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. Graphs with more than one type of edges are not allowed.

Returns

Return degree embedding vectors of shape \((N, d)\), where \(N\) is the number of nodes in the input graph and \(d\) is embedding_dim.

Return type

Tensor