DegreeEncoder

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

基类: Module

度编码器（Degree Encoder），如论文 Do Transformers Really Perform Bad for Graph Representation? 中介绍的那样。

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

参数:

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

参数:

degrees (Tensor) – 如果 direction 是 both，它应该是经过零填充的批处理图的入度和出度的堆叠，形状为 \((2, B, N)\) 的张量。否则，它应该是经过零填充的批处理图的入度或出度，形状为 \((B, N)\) 的张量，其中 \(B\) 是批处理图的批大小，\(N\) 是节点的最大数量。

返回值:

返回形状为 \((B, N, d)\) 的度嵌入向量，其中 \(d\) 是 embedding_dim。

返回类型:

Tensor