transformer

Transformer building blocks for point-cloud backbones and detectors.

Attention and TransformerBlock are the plain ViT-style multi-head self-attention and pre-norm residual block over dense token sequences of shape \((B, N, C)\). TransformerEncoderLayer and TransformerDecoderLayer are the DETR-style layers with positional embeddings added to the attention inputs, over sequence-first \((S, B, C)\) tokens.

Classes:

Attention –

Multi-head self-attention over a dense token sequence.
TransformerBlock –

Pre-norm transformer block: residual multi-head attention then a residual MLP.
TransformerEncoderLayer –

DETR-style transformer encoder layer: self-attention with positional embeddings added to the query and key,
TransformerDecoderLayer –

DETR-style transformer decoder layer: self-attention over the queries, cross-attention onto the memory and a

Attention [source] ¶

Attention(
    dim: int,
    num_heads: int = 8,
    qkv_bias: bool = False,
    qk_scale: Optional[float] = None,
    attn_dropout: float = 0.0,
    proj_dropout: float = 0.0,
)

Bases: Module

Multi-head self-attention over a dense token sequence.

Computes scaled dot-product attention with a single fused qkv projection and an output proj. An optional additive mask is added to the pre-softmax attention logits, which supports local / windowed attention (masked positions get a large negative bias).

Parameters:

dim (int) –

Token dimension \(C\). Must be divisible by num_heads.
num_heads (int, default: 8 ) –

Number of attention heads \(h\).
qkv_bias (bool, default: False ) –

Whether the fused query/key/value projection uses a bias.
qk_scale (Optional[float], default: None ) –

Override for the \(1/\sqrt{d_\text{head}}\) logit scale. Defaults to \(d_\text{head}^{-1/2}\) when None.
attn_dropout (float, default: 0.0 ) –

Dropout applied to the attention weights.
proj_dropout (float, default: 0.0 ) –

Dropout applied to the output projection.

Shape

Input: \((B, N, C)\) tokens and an optional mask broadcastable to \((B, h, N, N)\).
Output: \((B, N, C)\).

Example

import torch
from torch_pointcloud.layers import Attention

attn = Attention(384, num_heads=6)
x = torch.randn(2, 64, 384)
y = attn(x)
print(y.shape)

TransformerBlock [source] ¶

TransformerBlock(
    dim: int,
    num_heads: int,
    mlp_ratio: float = 4.0,
    qkv_bias: bool = False,
    qk_scale: Optional[float] = None,
    dropout: float = 0.0,
    attn_dropout: float = 0.0,
    drop_path: float = 0.0,
    act: Union[str, Callable, None] = "gelu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = LayerNorm,
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Pre-norm transformer block: residual multi-head attention then a residual MLP.

Applies \(x \leftarrow x + \text{DropPath}(\text{Attn}(\text{Norm}(x)))\) followed by \(x \leftarrow x + \text{DropPath}(\text{MLP}(\text{Norm}(x)))\). The feed-forward is a plain-last torch_geometric.nn.MLP of hidden size \(\lfloor C \cdot \text{mlp\_ratio} \rfloor\), so activation and dropout are configurable through the resolver API.

Parameters:

dim (int) –

Token dimension \(C\).
num_heads (int) –

Number of attention heads.
mlp_ratio (float, default: 4.0 ) –

Hidden-to-input ratio of the feed-forward MLP.
qkv_bias (bool, default: False ) –

Whether the attention qkv projection uses a bias.
qk_scale (Optional[float], default: None ) –

Override for the attention logit scale.
dropout (float, default: 0.0 ) –

Dropout used in the MLP and the attention output projection.
attn_dropout (float, default: 0.0 ) –

Dropout applied to the attention weights.
drop_path (float, default: 0.0 ) –

Stochastic-depth rate for the two residual branches.
act (Union[str, Callable, None], default: 'gelu' ) –

Activation for the feed-forward MLP.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra arguments for the activation.
norm (Union[str, Callable, None], default: LayerNorm ) –

Normalization applied before attention and before the MLP.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra arguments for the normalization.

Shape

Input: \((B, N, C)\) tokens and an optional mask broadcastable to \((B, h, N, N)\).
Output: \((B, N, C)\).

Example

import torch
from torch_pointcloud.layers import TransformerBlock

block = TransformerBlock(384, num_heads=6, drop_path=0.1)
x = torch.randn(2, 64, 384)
y = block(x)
print(y.shape)

TransformerEncoderLayer [source] ¶

TransformerEncoderLayer(
    embed_dim: int,
    num_heads: int,
    mlp_dim: int,
    dropout: float,
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_first: bool = True,
)

Bases: Module

DETR-style transformer encoder layer: self-attention with positional embeddings added to the query and key, then a feed-forward block, each wrapped in a residual.

Tokens are sequence-first, \((S, B, C)\), as nn.MultiheadAttention expects. With norm_first the layer norms are applied to the input of each sub-block (pre-norm); otherwise after each residual sum (post-norm).

Parameters:

embed_dim (int) –

Token embedding dimension.
num_heads (int) –

Number of attention heads.
mlp_dim (int) –

Hidden width of the feed-forward block.
dropout (float) –

Dropout probability.
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for the feed-forward block.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm_first (bool, default: True ) –

Whether to normalize the input of each sub-block rather than its residual sum.

Shape

Input: \((S, B, C)\) tokens, an optional \((S, S)\) attention mask, \((B, S)\) key padding mask and \((S, B, C)\) positions Output: \((S, B, C)\) tokens

TransformerDecoderLayer [source] ¶

TransformerDecoderLayer(
    embed_dim: int,
    num_heads: int,
    mlp_dim: int,
    dropout: float,
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm_first: bool = True,
    pos_in_value: bool = False,
)

Bases: Module

DETR-style transformer decoder layer: self-attention over the queries, cross-attention onto the memory and a feed-forward block, each wrapped in a residual.

Query positions are added to the self-attention query and key and to the cross-attention query; memory positions are added to the cross-attention key. With pos_in_value both are added to the attention values as well. Tokens are sequence-first, \((S, B, C)\). With norm_first the layer norms are applied to the input of each sub-block (pre-norm); otherwise after each residual sum (post-norm).

Parameters:

embed_dim (int) –

Token embedding dimension.
num_heads (int) –

Number of attention heads.
mlp_dim (int) –

Hidden width of the feed-forward block.
dropout (float) –

Dropout probability.
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for the feed-forward block.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm_first (bool, default: True ) –

Whether to normalize the input of each sub-block rather than its residual sum.
pos_in_value (bool, default: False ) –

Whether the positional embeddings are also added to the attention values.

Shape

Input: \((T, B, C)\) queries, \((S, B, C)\) memory, \((S, B, C)\) memory positions, \((T, B, C)\) query positions and an optional \((B, S)\) memory key padding mask Output: \((T, B, C)\) queries