Pooling Layers
The k3_node.layers.pool module provides global readouts, hierarchical graph coarsening, cluster-based pooling, and spatial graph construction operators.
Global Graph Readout
global_add_pool
k3_node.layers.pool.global_add_pool(x, batch=None, size=None)
Returns batch-wise graph-level-outputs by adding node features across the node dimension.
global_mean_pool
k3_node.layers.pool.global_mean_pool(x, batch=None, size=None)
Returns batch-wise graph-level-outputs by averaging node features across the node dimension.
global_max_pool
k3_node.layers.pool.global_max_pool(x, batch=None, size=None)
Returns batch-wise graph-level-outputs by taking the channel-wise maximum across the node dimension.
Hierarchical Node Pooling
TopKPooling
k3_node.layers.pool.TopKPooling
Bases: Layer
:math:\mathrm{top}_k pooling operator from the "Graph U-Nets"
<https://arxiv.org/abs/1905.05178>, "Towards Sparse
Hierarchical Graph Classifiers" <https://arxiv.org/abs/1811.01287>
and "Understanding Attention and Generalization in Graph Neural
Networks" <https://arxiv.org/abs/1905.02850>_ papers.
call(x, edge_index, edge_attr=None, batch=None, attn=None)
Forward pass.
reset_parameters()
Resets all learnable parameters of the module.
SAGPooling
k3_node.layers.pool.SAGPooling
Bases: Layer
The self-attention pooling operator from the "Self-Attention Graph
Pooling" <https://arxiv.org/abs/1904.08082> and "Understanding
Attention and Generalization in Graph Neural Networks"
<https://arxiv.org/abs/1905.02850> papers.
call(x, edge_index, edge_attr=None, batch=None, attn=None)
Forward pass.
reset_parameters()
Resets all learnable parameters of the module.
EdgePooling
k3_node.layers.pool.EdgePooling
Bases: Layer
The edge pooling operator from the "Towards Graph Pooling by Edge
Contraction" <https://graphreason.github.io/papers/17.pdf> and
"Edge Contraction Pooling for Graph Neural Networks"
<https://arxiv.org/abs/1905.10990> papers.
call(x, edge_index, batch, training=False)
Forward pass.
unpool(x, unpool_info)
Unpools a previous edge pooling step.
ASAPooling
k3_node.layers.pool.ASAPooling
Bases: Layer
The Adaptive Structure Aware Pooling operator from the
"ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical
Graph Representations" <https://arxiv.org/abs/1911.07979>_ paper.
call(x, edge_index, edge_weight=None, batch=None, training=False)
Forward pass.
PANPooling
k3_node.layers.pool.PANPooling
Bases: Layer
The path integral based pooling operator from the
"Path Integral Based Convolution and Pooling for Graph Neural Networks"
<https://arxiv.org/abs/2006.16811>_ paper.
call(x, M, batch=None)
Forward pass.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
x
|
Node feature matrix. |
required | |
M
|
MET matrix, either a tuple/list (edge_index, edge_weight) or an object
with |
required | |
batch
|
Optional[any]
|
Batch vector. |
None
|
MemPooling
k3_node.layers.pool.MemPooling
Bases: Layer
Memory based pooling layer from "Memory-Based Graph Networks"
<https://arxiv.org/abs/2002.09518>_ paper.
call(x, batch=None, mask=None, max_num_nodes=None, batch_size=None)
Forward pass.
kl_loss(S)
staticmethod
The additional KL divergence-based loss.
ClusterPooling
k3_node.layers.pool.ClusterPooling
Bases: Layer
The cluster pooling operator from the "Edge-Based Graph Component
Pooling" <https://arxiv.org/abs/2409.11856>_ paper.
call(x, edge_index, batch, training=False)
Forward pass.
Spatial & Neighborhood Pooling
avg_pool_neighbor_x
k3_node.layers.pool.avg_pool_neighbor_x(data, edge_index=None, flow='source_to_target')
Average-pools neighboring node features.
max_pool_neighbor_x
k3_node.layers.pool.max_pool_neighbor_x(data, edge_index=None, flow='source_to_target')
Max-pools neighboring node features.
voxel_grid
k3_node.layers.pool.voxel_grid
voxel_grid(pos, size, batch=None, start=None, end=None)
Voxel grid pooling that clusters points within the same voxel.
fps
k3_node.layers.pool.fps(x, batch=None, ratio=0.5, random_start=True, batch_size=None)
Farthest Point Sampling algorithm.
graclus
k3_node.layers.pool.graclus
graclus(edge_index, weight=None, num_nodes=None)
A greedy clustering algorithm of picking an unmarked vertex and matching it with one of its unmarked neighbors that maximizes its edge weight.
Graph Construction
radius_graph
k3_node.layers.pool.radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32, flow='source_to_target', num_workers=1, batch_size=None)
Computes graph edges to all points within a given distance r.
knn_graph
k3_node.layers.pool.knn_graph(x, k, batch=None, loop=False, flow='source_to_target', cosine=False, num_workers=1, batch_size=None)
Computes graph edges to the nearest k points.