torch_geometric.nn.aggr.Set2Set

class Set2Set(in_channels: int, processing_steps: int, **kwargs)

Bases: Aggregation

基于迭代内容注意力的Set2Set聚合操作符，如“顺序重要：集合的序列到序列”论文中所述。

\[ \begin{align}\begin{aligned}\mathbf{q}_t &= \mathrm{LSTM}(\mathbf{q}^{*}_{t-1})\\\alpha_{i,t} &= \mathrm{softmax}(\mathbf{x}_i \cdot \mathbf{q}_t)\\\mathbf{r}_t &= \sum_{i=1}^N \alpha_{i,t} \mathbf{x}_i\\\mathbf{q}^{*}_t &= \mathbf{q}_t \, \Vert \, \mathbf{r}_t,\end{aligned}\end{align} \]

其中 \(\mathbf{q}^{*}_T\) 定义了该层的输出，其维度是输入的两倍。

Parameters:

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

• processing_steps (int) – 迭代次数 \(T\).

• **kwargs (optional) – Additional arguments of torch.nn.LSTM.

reset_parameters()

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

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

前向传播。

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 – (int, optional): The maximum number of elements within a single aggregation group. (default: None)

Return type:

Tensor