torch_geometric.nn.aggr.GraphMultisetTransformer

class GraphMultisetTransformer(channels: int, k: int, num_encoder_blocks: int = 1, heads: int = 1, layer_norm: bool = False, dropout: float = 0.0)

基础类：Aggregation

来自“Accurate Learning of Graph Representations with Graph Multiset Pooling”论文的图多集变换器池化操作符。

GraphMultisetTransformer 通过基于注意力的池化将元素聚合为 \(k\) 个代表性元素，通过 num_encoder_blocks 自注意力块计算它们之间的交互， 最后通过基于注意力的池化将代表性元素聚合成一个单一的簇。

注意

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

Parameters:

• channels (int) – 每个输入样本的大小。

• k (int) – 池化后的\(k\)个代表节点的数量。

• num_encoder_blocks (int, optional) – 两个池化块之间的集合注意力块（SABs）的数量。（默认值：1）

• heads (int, optional) – 多头注意力机制的数量。 (默认: 1)

• norm (str, optional) – 如果设置为 True，将应用层归一化。(默认值: False)

• dropout (float, optional) – 注意力权重的丢弃概率。 (默认值: 0)

reset_parameters()

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

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

前向传播。

Parameters:

• x (torch.Tensor) – 源张量。

• index (torch.Tensor, optional) – 用于应用聚合的元素的索引。 必须定义 index 或 ptr 中的一个。 (默认值: None)

• ptr (torch.Tensor, optional) – 如果提供，将基于CSR表示中的排序输入计算聚合。 必须定义index或ptr中的一个。 (默认: None)

• dim_size (int, optional) – 聚合后输出张量在维度 dim 的大小。（默认值：None）

• dim (int, optional) – 聚合的维度。 (default: -2)

• max_num_elements (Optional[int], 默认值: None) – (int, 可选): 单个聚合组中的最大元素数量。(默认值: None)

Return type:

Tensor