transformer

Standard pre-norm transformer building blocks for point-cloud backbones.

The Attention and TransformerBlock here are the plain ViT-style multi-head self-attention and residual block. They operate on dense token sequences of shape \((B, N, C)\).

Classes:

Attention –

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

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

Attention ¶

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 ¶

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)