torch_geometric.nn.aggr.SortAggregation

class SortAggregation(k: int)[source]

Bases: Aggregation

来自“An End-to-End Deep Learning Architecture for Graph Classification”论文的池化操作符， 其中节点特征根据它们的最后一个特征通道按降序排序。前\(k\)个节点形成该层的输出。

注意

SortAggregation 需要排序的索引 index 作为输入。 具体来说，如果你将此聚合作为 MessagePassing 的一部分使用，请确保 edge_index 按目标节点排序，可以通过手动 使用 sort_edge_index() 排序边索引 或调用 torch_geometric.data.Data.sort() 来实现。

Parameters:

k (int) – 每个图要保留的节点数。

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