torch_geometric.nn.aggr.Aggregation

class Aggregation[source]

Bases: Module

用于实现自定义聚合的抽象基类。

聚合可以通过一个index向量来执行,该向量定义了从输入元素到它们在输出中位置的映射:


https://raw.githubusercontent.com/rusty1s/pytorch_scatter/master/docs/source/_figures/add.svg?sanitize=true

值得注意的是,index 不需要排序(对于大多数聚合操作符):

# Feature matrix holding 10 elements with 64 features each:
x = torch.randn(10, 64)

# Assign each element to one of three sets:
index = torch.tensor([0, 0, 1, 0, 2, 0, 2, 1, 0, 2])

output = aggr(x, index)  #  Output shape: [3, 64]

或者,可以通过一个称为ptr的“压缩”索引向量来实现聚合。在这里,同一集合中的元素需要在输入中分组在一起,而ptr定义了它们的边界:

# Feature matrix holding 10 elements with 64 features each:
x = torch.randn(10, 64)

# Define the boundary indices for three sets:
ptr = torch.tensor([0, 4, 7, 10])

output = aggr(x, ptr=ptr)  #  Output shape: [3, 64]

请注意,至少需要定义indexptr中的一个。

Shapes:
  • 输入: 节点特征 \((*, |\mathcal{V}|, F_{in})\) 或边特征 \((*, |\mathcal{E}|, F_{in})\), 索引向量 \((|\mathcal{V}|)\)\((|\mathcal{E}|)\),

  • 输出: 图特征 \((*, |\mathcal{G}|, F_{out})\) 或 节点特征 \((*, |\mathcal{V}|, F_{out})\)

forward(x: Tensor, index: Optional[Tensor] = None, ptr: Optional[Tensor] = None, dim_size: Optional[int] = None, dim: int = -2, max_num_elements: Optional[int] = None) Tensor[source]

前向传播。

Parameters:
  • x (torch.Tensor) – The source tensor.

  • index (torch.Tensor, optional) – The indices of elements for applying the aggregation. One of index or ptr must be defined. (default: None)

  • ptr (torch.Tensor, optional) – If given, computes the aggregation based on sorted inputs in CSR representation. One of index or ptr must be defined. (default: None)

  • dim_size (int, optional) – The size of the output tensor at dimension dim after aggregation. (default: None)

  • dim (int, optional) – The dimension in which to aggregate. (default: -2)

  • max_num_elements (Optional[int], default: None) – (int, optional): The maximum number of elements within a single aggregation group. (default: None)

Return type:

Tensor

reset_parameters()[source]

重置模块的所有可学习参数。