torch_geometric.nn.aggr.EquilibriumAggregation

class EquilibriumAggregation(in_channels: int, out_channels: int, num_layers: List[int], grad_iter: int = 5, lamb: float = 0.1)

Bases: Aggregation

来自“Equilibrium Aggregation: Encoding Sets via Optimization”论文的均衡聚合层。

该层的输出 \(\mathbf{y}\) 是通过一个势函数 \(F(\mathbf{x}, \mathbf{y})\)、一个正则化项 \(R(\mathbf{y})\) 以及条件隐式定义的。

\[\mathbf{y} = \min_\mathbf{y} R(\mathbf{y}) + \sum_{i} F(\mathbf{x}_i, \mathbf{y}).\]

给定的实现使用了一个类似ResNet的模型作为潜在函数，并使用了一个简单的\(L_2\)范数\(R(\mathbf{y}) = \textrm{softplus}(\lambda) \cdot {\| \mathbf{y} \|}^2_2\)作为正则化器，其中\(\lambda\)是可学习的权重。

Parameters:

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

• out_channels (int) – Size of each output sample.

• num_layers (List[int) – 潜在函数中隐藏通道的列表。

• grad_iter (int) – 内部梯度下降的步数。（默认值：5）

• lamb (float) – 初始正则化常数。 (默认值: 0.1)

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