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_scatter.scatter(reduce="max").

  • MinPool –

    Per-segment min pooling over a packed batch, via torch_scatter.scatter(reduce="min").

  • MeanPool –

    Per-segment mean pooling over a packed batch, via torch_scatter.scatter(reduce="mean").

  • MulPool –

    Per-segment product pooling over a packed batch, via torch_scatter.scatter(reduce="mul").

  • SumPool –

    Per-segment sum pooling over a packed batch, via torch_scatter.scatter(reduce="sum").

  • SoftmaxPool –

    Per-segment softmax pooling, delegating reduce="softmax" to torch_scatter.scatter.

  • LogSoftmaxPool –

    Per-segment log-softmax pooling, delegating reduce="log_softmax" to torch_scatter.scatter.

  • 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.

  • create_adaptive_pool –

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

MaxPool 

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

Bases: Module

Per-segment max pooling over a packed batch, via torch_scatter.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 

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

Bases: Module

Per-segment min pooling over a packed batch, via torch_scatter.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 

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

Bases: Module

Per-segment mean pooling over a packed batch, via torch_scatter.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 

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

Bases: Module

Per-segment product pooling over a packed batch, via torch_scatter.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 

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

Bases: Module

Per-segment sum pooling over a packed batch, via torch_scatter.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.

SoftmaxPool 

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

Bases: Module

Per-segment softmax pooling, delegating reduce="softmax" to torch_scatter.scatter.

Warning

torch_scatter.scatter only accepts sum / mean / min / max / mul reductions, so calling this module raises ValueError with current torch_scatter releases.

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.

LogSoftmaxPool 

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

Bases: Module

Per-segment log-softmax pooling, delegating reduce="log_softmax" to torch_scatter.scatter.

Warning

torch_scatter.scatter only accepts sum / mean / min / max / mul reductions, so calling this module raises ValueError with current torch_scatter releases.

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.

CatPool 

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 

create_pool(
    name: PoolLike, *args: Any, **kwargs: Any
) -> Module

Resolve a packed-batch pooling module from a name, class, or instance.

Parameters:

  • name (PoolLike) –

    Pool name ("max", "min", "mean", "mul", "sum", "softmax", "log_softmax"), a module class, or an existing instance (returned as-is).

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

    Positional arguments forwarded to the pool constructor.

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

    Keyword arguments forwarded to the pool constructor (dim, dim_size).

Returns:

  • Module –

    The instantiated pooling module.

create_adaptive_pool 

create_adaptive_pool(
    name: AdaptivePoolLike, *args: Any, **kwargs: Any
) -> Module

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

Parameters:

  • name (AdaptivePoolLike) –

    Pool name ("mean", "max"), a module class, or an existing instance (returned as-is).

  • *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:

  • Module –

    The instantiated pooling module.