from torch_scatter import scatter
import torch_geometric.graphgym.register as register
[docs]def global_add_pool(x, batch, size=None):
"""
Globally pool node embeddings into graph embeddings, via elementwise sum.
Pooling function takes in node embedding [num_nodes x emb_dim] and
batch (indices) and outputs graph embedding [num_graphs x emb_dim].
Args:
x (torch.tensor): Input node embeddings
batch (torch.tensor): Batch tensor that indicates which node
belongs to which graph
size (optional): Total number of graphs. Can be auto-inferred.
Returns: Pooled graph embeddings
"""
size = batch.max().item() + 1 if size is None else size
return scatter(x, batch, dim=0, dim_size=size, reduce='add')
[docs]def global_mean_pool(x, batch, size=None):
"""
Globally pool node embeddings into graph embeddings, via elementwise mean.
Pooling function takes in node embedding [num_nodes x emb_dim] and
batch (indices) and outputs graph embedding [num_graphs x emb_dim].
Args:
x (torch.tensor): Input node embeddings
batch (torch.tensor): Batch tensor that indicates which node
belongs to which graph
size (optional): Total number of graphs. Can be auto-inferred.
Returns: Pooled graph embeddings
"""
size = batch.max().item() + 1 if size is None else size
return scatter(x, batch, dim=0, dim_size=size, reduce='mean')
[docs]def global_max_pool(x, batch, size=None):
"""
Globally pool node embeddings into graph embeddings, via elementwise max.
Pooling function takes in node embedding [num_nodes x emb_dim] and
batch (indices) and outputs graph embedding [num_graphs x emb_dim].
Args:
x (torch.tensor): Input node embeddings
batch (torch.tensor): Batch tensor that indicates which node
belongs to which graph
size (optional): Total number of graphs. Can be auto-inferred.
Returns: Pooled graph embeddings
"""
size = batch.max().item() + 1 if size is None else size
return scatter(x, batch, dim=0, dim_size=size, reduce='max')
pooling_dict = {
'add': global_add_pool,
'mean': global_mean_pool,
'max': global_max_pool
}
register.pooling_dict = {**register.pooling_dict, **pooling_dict}