pools
Per-segment pooling modules over packed batches and the create_pool / create_adaptive_pool factories.
Classes:
-
MaxPoolβPer-segment max pooling over a packed batch, via
torch_geometric.utils.scatter(reduce="max"). -
MinPoolβPer-segment min pooling over a packed batch, via
torch_geometric.utils.scatter(reduce="min"). -
MeanPoolβPer-segment mean pooling over a packed batch, via
torch_geometric.utils.scatter(reduce="mean"). -
MulPoolβPer-segment product pooling over a packed batch, via
torch_geometric.utils.scatter(reduce="mul"). -
SumPoolβPer-segment sum pooling over a packed batch, via
torch_geometric.utils.scatter(reduce="sum"). -
CatPoolβRuns several pools on the same input and concatenates their outputs along the feature dim.
Functions:
-
create_poolβResolve a packed-batch pooling module from a name, class, or instance, returning
Noneforname=None. -
create_adaptive_poolβResolve a dense adaptive pooling module (
nn.AdaptiveAvgPool1d/nn.AdaptiveMaxPool1d), returningNone
MaxPool [source] ΒΆ
Bases: Module
Per-segment max pooling over a packed batch, via torch_geometric.utils.scatter(reduce="max").
Parameters:
-
dim(int, default:0) βDimension along which to pool.
-
dim_size(Optional[int], default:None) βNumber of output segments \(B\).
Noneinfers it from the segment index.
Shape
Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.
MinPool [source] ΒΆ
Bases: Module
Per-segment min pooling over a packed batch, via torch_geometric.utils.scatter(reduce="min").
Parameters:
-
dim(int, default:0) βDimension along which to pool.
-
dim_size(Optional[int], default:None) βNumber of output segments \(B\).
Noneinfers it from the segment index.
Shape
Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.
MeanPool [source] ΒΆ
Bases: Module
Per-segment mean pooling over a packed batch, via torch_geometric.utils.scatter(reduce="mean").
Parameters:
-
dim(int, default:0) βDimension along which to pool.
-
dim_size(Optional[int], default:None) βNumber of output segments \(B\).
Noneinfers it from the segment index.
Shape
Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.
MulPool [source] ΒΆ
Bases: Module
Per-segment product pooling over a packed batch, via torch_geometric.utils.scatter(reduce="mul").
Parameters:
-
dim(int, default:0) βDimension along which to pool.
-
dim_size(Optional[int], default:None) βNumber of output segments \(B\).
Noneinfers it from the segment index.
Shape
Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.
SumPool [source] ΒΆ
Bases: Module
Per-segment sum pooling over a packed batch, via torch_geometric.utils.scatter(reduce="sum").
Parameters:
-
dim(int, default:0) βDimension along which to pool.
-
dim_size(Optional[int], default:None) βNumber of output segments \(B\).
Noneinfers it from the segment index.
Shape
Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.
CatPool [source] ΒΆ
CatPool(
pools: Sequence[PoolLike] = ("max", "mean"),
dim: int = 0,
dim_size: Optional[int] = None,
)
Bases: Module
Runs several pools on the same input and concatenates their outputs along the feature dim.
Parameters:
-
pools(Sequence[PoolLike], default:('max', 'mean')) βPools to combine, each resolved by
create_pool(a name, class, or instance). -
dim(int, default:0) βDimension along which each pool reduces.
-
dim_size(Optional[int], default:None) βNumber of output segments \(B\).
Noneinfers it from the segment index.
Shape
Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C \cdot P)\) where \(P\) is the number of pools.
Example
Attributes:
-
num_pools(int) βNumber of pools \(P\) concatenated, i.e. the feature multiplier.
num_pools property ΒΆ
Number of pools \(P\) concatenated, i.e. the feature multiplier.
create_pool [source] ΒΆ
Resolve a packed-batch pooling module from a name, class, or instance, returning None for name=None.
Parameters:
-
name(Optional[PoolLike]) βPool name (
"max","min","mean","mul","sum"), a module class, an existing instance (returned as-is), orNone. -
*args(Any, default:()) βPositional arguments forwarded to the pool constructor.
-
**kwargs(Any, default:{}) βKeyword arguments forwarded to the pool constructor (
dim,dim_size).
Returns:
-
Optional[Module]βThe instantiated pooling module, or
Noneifname is None.
create_adaptive_pool [source] ΒΆ
create_adaptive_pool(
name: Optional[AdaptivePoolLike],
*args: Any,
**kwargs: Any,
) -> Optional[Module]
Resolve a dense adaptive pooling module (nn.AdaptiveAvgPool1d / nn.AdaptiveMaxPool1d), returning None for name=None.
Parameters:
-
name(Optional[AdaptivePoolLike]) βPool name (
"mean","max"), a module class, an existing instance (returned as-is), orNone. -
*args(Any, default:()) βPositional arguments forwarded to the pool constructor.
-
**kwargs(Any, default:{}) βKeyword arguments forwarded to the pool constructor.
output_sizedefaults to1(global pooling).
Returns:
-
Optional[Module]βThe instantiated pooling module, or
Noneifname is None.