DegreeEncoder

class dgl.nn.pytorch.gt.DegreeEncoder(max_degree, embedding_dim, direction='both')[source]

Bases: Module

Degree Encoder，如Do Transformers Really Perform Bad for Graph Representation?中介绍

该模块是一个可学习的度嵌入模块。

Parameters:

max_degree (int) – 要编码的度数的上限。每个度数将被限制在范围 [0, max_degree] 内。
embedding_dim (int) – 嵌入向量的输出维度。
direction (str, optional) – 要编码的方向的度数，从 in, out 和 both 中选择。 both 编码来自两个方向的度数并输出它们的和。默认值 : both.

示例

>>> 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])

forward(degrees)[source]

Parameters:: degrees (Tensor) – 如果 direction 是 both，它应该以零填充的方式堆叠批处理图的入度和出度，形状为 \((2, B, N)\) 的张量。否则，它应该以零填充的方式堆叠批处理图的入度或出度，形状为 \((B, N)\) 的张量，其中 \(B\) 是批处理图的批量大小， \(N\) 是最大节点数。
Returns:: 返回形状为 \((B, N, d)\) 的度嵌入向量，其中 \(d\) 是 embedding_dim。
Return type:: 张量