torch_geometric.nn.aggr.SetTransformerAggregation

class SetTransformerAggregation(channels: int, num_seed_points: int = 1, num_encoder_blocks: int = 1, num_decoder_blocks: int = 1, heads: int = 1, concat: bool = True, layer_norm: bool = False, dropout: float = 0.0)[source]

Bases: Aggregation

执行“Set Transformer”聚合，其中要聚合的元素通过多头注意力块进行处理，如“具有自适应读数的图神经网络”论文中所述。

注意

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

Parameters:

  • channels (int) – Size of each input sample.

  • num_seed_points (int, optional) – 种子点的数量。 (默认: 1)

  • num_encoder_blocks (int, optional) – 编码器中Set Attention Blocks (SABs)的数量。(默认: 1).

  • num_decoder_blocks (int, optional) – 解码器中Set Attention Blocks (SABs)的数量。(默认: 1).

  • heads (int, optional) – Number of multi-head-attentions. (default: 1)

  • concat (bool, 可选) – 如果设置为 False，种子嵌入将被平均而不是连接。（默认值：True）

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

  • dropout (float, optional) – Dropout probability of attention weights. (default: 0)

reset_parameters()[source]

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

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