Skip to content

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 None for name=None.

  • create_adaptive_pool –

    Resolve a dense adaptive pooling module (nn.AdaptiveAvgPool1d / nn.AdaptiveMaxPool1d), returning None

MaxPool [source] ΒΆ

MaxPool(dim: int = 0, dim_size: Optional[int] = None)

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\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.

MinPool [source] ΒΆ

MinPool(dim: int = 0, dim_size: Optional[int] = None)

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\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.

MeanPool [source] ΒΆ

MeanPool(dim: int = 0, dim_size: Optional[int] = None)

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\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.

MulPool [source] ΒΆ

MulPool(dim: int = 0, dim_size: Optional[int] = None)

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\). None infers it from the segment index.

Shape

Input: \((N, C)\) features x and a \((N,)\) segment index batch. Output: \((B, C)\) pooled features.

SumPool [source] ΒΆ

SumPool(dim: int = 0, dim_size: Optional[int] = None)

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\). None infers 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\). None infers 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
import torch
from torch_pointcloud.layers import CatPool

pool = CatPool(pools=("max", "mean"))
x = torch.randn(6, 4)
batch = torch.tensor([0, 0, 0, 1, 1, 1])
out = pool(x, batch)  # (2, 8)

Attributes:

  • num_pools (int) –

    Number of pools \(P\) concatenated, i.e. the feature multiplier.

num_pools property ΒΆ

num_pools: int

Number of pools \(P\) concatenated, i.e. the feature multiplier.

create_pool [source] ΒΆ

create_pool(name: None, *args: Any, **kwargs: Any) -> None
create_pool(
    name: PoolLike, *args: Any, **kwargs: Any
) -> Module
create_pool(
    name: Optional[PoolLike], *args: Any, **kwargs: Any
) -> Optional[Module]

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), or None.

  • *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 None if name is None.

create_adaptive_pool [source] ΒΆ

create_adaptive_pool(
    name: None, *args: Any, **kwargs: Any
) -> None
create_adaptive_pool(
    name: AdaptivePoolLike, *args: Any, **kwargs: Any
) -> Module
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), or None.

  • *args (Any, default: () ) –

    Positional arguments forwarded to the pool constructor.

  • **kwargs (Any, default: {} ) –

    Keyword arguments forwarded to the pool constructor. output_size defaults to 1 (global pooling).

Returns:

  • Optional[Module] –

    The instantiated pooling module, or None if name is None.