"""Point Transformer V3 with 3D RoPE, bottleneck CPE, and QK-Norm. Mostly based on Utonia's
(arXiv:2603.03283) PTv3 implementation:
https://github.com/Pointcept/Pointcept/blob/main/pointcept/models/point_transformer_v3/point_transformer_v3m3_utonia.py
Registered as both PT-v3m4 and PT-v3m8.
"""
from spconv import constants
constants.SPCONV_ALLOW_TF32 = True
import spconv.pytorch as spconv
import torch
import torch.nn as nn
import torch_scatter
import pointops
from addict import Dict
from timm.layers import DropPath
from torch.nn.init import trunc_normal_
from pimm.models.builder import MODELS
from pimm.models.modules import PointModule, PointSequential
from pimm.models.utils.misc import offset2bincount
from pimm.models.utils.structure import Point
try:
from flash_attn import flash_attn_varlen_func
except ImportError:
flash_attn_varlen_func = None
class Point3DRoPE(nn.Module):
"""3D rotary position embedding for serialized point cloud attention.
Splits each head into 3 equal sub-vectors (x, y, z) and applies 1D RoPE
per axis. head_dim must be divisible by 6 (3 axes * 2 for rotation pairs).
Optional coordinate augmentation during training (DINOv3-style):
- jitter: axis-wise multiplicative perturbation, j = exp(U(-log γ, log γ)^3)
- rescale: isotropic multiplicative perturbation, r = exp(U(-log η, log η))
"""
def __init__(self, head_dim, base=10000, jitter_degree=None, rescale_degree=None):
super().__init__()
assert (
head_dim % 6 == 0
), f"head_dim must be divisible by 6 for 3D RoPE, got {head_dim}"
self.head_dim = head_dim
self.chunk_dim = head_dim // 3
self.base = base
self.jitter_degree = jitter_degree
self.rescale_degree = rescale_degree
inv_freq = 1.0 / (
self.base ** (torch.arange(0, self.chunk_dim, 2).float() / self.chunk_dim)
)
self.register_buffer("inv_freq", inv_freq)
@torch.no_grad()
def _augment_coords(self, xyz):
if not self.training:
return xyz
if self.jitter_degree is not None and self.jitter_degree > 1:
log_gamma = torch.tensor(
self.jitter_degree, device=xyz.device, dtype=xyz.dtype
).log()
eps_j = torch.empty(3, device=xyz.device, dtype=xyz.dtype).uniform_(
-log_gamma, log_gamma
)
xyz = xyz * eps_j.exp()
if self.rescale_degree is not None and self.rescale_degree > 1:
log_eta = torch.tensor(
self.rescale_degree, device=xyz.device, dtype=xyz.dtype
).log()
eps_s = torch.empty(1, device=xyz.device, dtype=xyz.dtype).uniform_(
-log_eta, log_eta
)
xyz = xyz * eps_s.exp()
return xyz
def _compute_cos_sin(self, xyz):
"""xyz: (..., 3) -> cos, sin each (..., head_dim)."""
x, y, z = xyz[..., 0:1], xyz[..., 1:2], xyz[..., 2:3]
emb_x = x * self.inv_freq
emb_y = y * self.inv_freq
emb_z = z * self.inv_freq
emb_x = torch.cat((emb_x, emb_x), dim=-1)
emb_y = torch.cat((emb_y, emb_y), dim=-1)
emb_z = torch.cat((emb_z, emb_z), dim=-1)
emb = torch.cat((emb_x, emb_y, emb_z), dim=-1)
return emb.cos(), emb.sin()
@staticmethod
def _rotate_half(x):
x1, x2 = x.chunk(2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
def forward(self, q, k, xyz):
"""Apply 3D RoPE to q and k.
q, k: (..., H, [K,] D) with head dim at position 1
xyz: (matching leading dims, [K,] 3)
"""
dtype = q.dtype
xyz = self._augment_coords(xyz)
cos, sin = self._compute_cos_sin(xyz)
# broadcast across heads (always dim 1)
cos = cos.unsqueeze(1)
sin = sin.unsqueeze(1)
q_chunks = q.split(self.chunk_dim, dim=-1)
k_chunks = k.split(self.chunk_dim, dim=-1)
cos_chunks = cos.split(self.chunk_dim, dim=-1)
sin_chunks = sin.split(self.chunk_dim, dim=-1)
q_out, k_out = [], []
for i in range(3):
q_out.append(
q_chunks[i] * cos_chunks[i]
+ self._rotate_half(q_chunks[i]) * sin_chunks[i]
)
k_out.append(
k_chunks[i] * cos_chunks[i]
+ self._rotate_half(k_chunks[i]) * sin_chunks[i]
)
return torch.cat(q_out, dim=-1).to(dtype), torch.cat(k_out, dim=-1).to(dtype)
class LayerScale(nn.Module):
def __init__(
self,
dim: int,
init_values: float = 1e-5,
inplace: bool = False,
) -> None:
super().__init__()
self.inplace = inplace
self.gamma = nn.Parameter(init_values * torch.ones(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
return x.mul_(self.gamma) if self.inplace else x * self.gamma
class RPE(torch.nn.Module):
def __init__(self, patch_size, num_heads):
super().__init__()
self.patch_size = patch_size
self.num_heads = num_heads
self.pos_bnd = int((4 * patch_size) ** (1 / 3) * 2)
self.rpe_num = 2 * self.pos_bnd + 1
self.rpe_table = torch.nn.Parameter(torch.zeros(3 * self.rpe_num, num_heads))
torch.nn.init.trunc_normal_(self.rpe_table, std=0.02)
def forward(self, coord):
idx = (
coord.clamp(-self.pos_bnd, self.pos_bnd) # clamp into bnd
+ self.pos_bnd # relative position to positive index
+ torch.arange(3, device=coord.device) * self.rpe_num # x, y, z stride
)
out = self.rpe_table.index_select(0, idx.reshape(-1))
out = out.view(idx.shape + (-1,)).sum(3)
out = out.permute(0, 3, 1, 2) # (N, K, K, H) -> (N, H, K, K)
return out
class SerializedAttention(PointModule):
def __init__(
self,
channels,
num_heads,
patch_size,
qk_norm=False,
qkv_bias=True,
qk_scale=None,
attn_drop=0.0,
proj_drop=0.0,
order_index=0,
enable_rpe=False,
enable_flash=True,
upcast_attention=True,
upcast_softmax=True,
rope_base=10,
rope_jitter=None,
rope_rescale=None,
):
super().__init__()
assert channels % num_heads == 0
self.channels = channels
self.num_heads = num_heads
head_dim = int(channels // num_heads)
self.scale = qk_scale or head_dim**-0.5
self.order_index = order_index
self.upcast_attention = upcast_attention
self.upcast_softmax = upcast_softmax
self.enable_rpe = enable_rpe
self.enable_flash = enable_flash
if enable_flash:
assert (
enable_rpe is False
), "Set enable_rpe to False when enable Flash Attention"
assert (
upcast_attention is False
), "Set upcast_attention to False when enable Flash Attention"
assert (
upcast_softmax is False
), "Set upcast_softmax to False when enable Flash Attention"
self.patch_size = patch_size
self.attn_drop = attn_drop
else:
self.patch_size_max = patch_size
self.patch_size = 0
self.attn_drop = torch.nn.Dropout(attn_drop)
if qk_norm:
self.q_norm = nn.LayerNorm(head_dim, eps=1e-6)
self.k_norm = nn.LayerNorm(head_dim, eps=1e-6)
else:
self.q_norm = None
self.k_norm = None
self.qkv = torch.nn.Linear(channels, channels * 3, bias=qkv_bias)
self.proj = torch.nn.Linear(channels, channels)
self.proj_drop = torch.nn.Dropout(proj_drop)
self.softmax = torch.nn.Softmax(dim=-1)
self.rpe = RPE(patch_size, num_heads) if self.enable_rpe else None
self.rope = Point3DRoPE(
head_dim,
rope_base,
jitter_degree=rope_jitter,
rescale_degree=rope_rescale,
)
@torch.no_grad()
def get_rel_pos(self, point, order):
K = self.patch_size
rel_pos_key = f"rel_pos_{self.order_index}"
if rel_pos_key not in point.keys():
grid_coord = point.grid_coord[order]
grid_coord = grid_coord.reshape(-1, K, 3)
point[rel_pos_key] = grid_coord.unsqueeze(2) - grid_coord.unsqueeze(1)
return point[rel_pos_key]
@torch.no_grad()
def get_padding_and_inverse(self, point):
pad_key = "pad"
unpad_key = "unpad"
cu_seqlens_key = "cu_seqlens_key"
if (
pad_key not in point.keys()
or unpad_key not in point.keys()
or cu_seqlens_key not in point.keys()
):
if self.patch_size == -1:
cu_seqlens = torch.cat([point.offset.new_zeros(1), point.offset]).int()
point[pad_key] = None
point[unpad_key] = None
point[cu_seqlens_key] = cu_seqlens
return point[pad_key], point[unpad_key], point[cu_seqlens_key]
offset = point.offset
bincount = offset2bincount(offset)
bincount_pad = (
torch.div(
bincount + self.patch_size - 1,
self.patch_size,
rounding_mode="trunc",
)
* self.patch_size
)
# only pad point when num of points larger than patch_size
mask_pad = bincount > self.patch_size
bincount_pad = ~mask_pad * bincount + mask_pad * bincount_pad
_offset = nn.functional.pad(offset, (1, 0))
_offset_pad = nn.functional.pad(torch.cumsum(bincount_pad, dim=0), (1, 0))
pad = torch.arange(_offset_pad[-1], device=offset.device)
unpad = torch.arange(_offset[-1], device=offset.device)
cu_seqlens = []
for i in range(len(offset)):
unpad[_offset[i] : _offset[i + 1]] += _offset_pad[i] - _offset[i]
if bincount[i] != bincount_pad[i]:
pad[
_offset_pad[i + 1]
- self.patch_size
+ (bincount[i] % self.patch_size) : _offset_pad[i + 1]
] = pad[
_offset_pad[i + 1]
- 2 * self.patch_size
+ (bincount[i] % self.patch_size) : _offset_pad[i + 1]
- self.patch_size
]
pad[_offset_pad[i] : _offset_pad[i + 1]] -= _offset_pad[i] - _offset[i]
cu_seqlens.append(
torch.arange(
_offset_pad[i],
_offset_pad[i + 1],
step=self.patch_size,
dtype=torch.int32,
device=offset.device,
)
)
point[pad_key] = pad
point[unpad_key] = unpad
point[cu_seqlens_key] = nn.functional.pad(
torch.concat(cu_seqlens), (0, 1), value=_offset_pad[-1]
)
return point[pad_key], point[unpad_key], point[cu_seqlens_key]
def forward(self, point, return_attn=False):
if not self.enable_flash:
self.patch_size = min(
offset2bincount(point.offset).min().tolist(), self.patch_size_max
)
H = self.num_heads
K = self.patch_size
C = self.channels
pad, unpad, cu_seqlens = self.get_padding_and_inverse(point)
max_seqlen = cu_seqlens[-1] if self.patch_size == -1 else self.patch_size
order = point.serialized_order[self.order_index]
if pad is not None:
order = order[pad]
inverse = unpad[point.serialized_inverse[self.order_index]]
qkv = self.qkv(point.feat)[order]
rope_coord = point.coord[order].clone()
if not self.enable_flash:
# (N', K, 3, H, C') => (3, N', H, K, C')
q, k, v = (
qkv.reshape(-1, K, 3, H, C // H).permute(2, 0, 3, 1, 4).unbind(dim=0)
)
# apply QK norm before RoPE!
if self.q_norm:
q = self.q_norm(q)
if self.k_norm:
k = self.k_norm(k)
# q: (P, H, K, D), rope_coord: (N_pad, 3) -> (P, K, 3)
q, k = self.rope(q, k, rope_coord.reshape(-1, K, 3))
if self.upcast_attention:
q = q.float()
k = k.float()
attn = (q * self.scale) @ k.transpose(-2, -1)
if self.enable_rpe:
attn = attn + self.rpe(self.get_rel_pos(point, order))
if self.upcast_softmax:
attn = attn.float()
attn = self.softmax(attn)
attn = self.attn_drop(attn).to(qkv.dtype)
feat = (attn @ v).transpose(1, 2).reshape(-1, C)
else:
# split q/k/v so we can apply RoPE before flash attention
qkv_bf = qkv.to(torch.bfloat16).reshape(-1, 3, H, C // H)
q, k, v = qkv_bf.unbind(dim=1) # each (N_pad, H, D)
if self.q_norm: # before rope!
q = self.q_norm(q)
if self.k_norm:
k = self.k_norm(k)
q, k = self.rope(q, k, rope_coord) # rope_coord: (N_pad, 3)
feat = flash_attn_varlen_func(
q, k, v,
cu_seqlens_q=cu_seqlens,
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
dropout_p=self.attn_drop if self.training else 0,
softmax_scale=self.scale,
).reshape(-1, C)
feat = feat.to(qkv.dtype)
if pad is not None:
feat = feat[inverse]
feat = self.proj(feat)
feat = self.proj_drop(feat)
point.feat = feat
return point
class MLP(nn.Module):
def __init__(
self,
in_channels,
hidden_channels=None,
out_channels=None,
act_layer=nn.GELU,
drop=0.0,
):
super().__init__()
out_channels = out_channels or in_channels
hidden_channels = hidden_channels or in_channels
self.fc1 = nn.Linear(in_channels, hidden_channels)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_channels, out_channels)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
class BottleneckCPE(PointModule):
"""CPE with an optional channel bottleneck around the sparse convolution.
Equal input and CPE widths skip the input projection, producing the former
full-width CPE path: sparse convolution, output projection, normalization.
"""
def __init__(self, channels, cpe_channels, cpe_indice_key, norm_layer=nn.LayerNorm):
super().__init__()
self.down = (
nn.Identity()
if channels == cpe_channels
else nn.Linear(channels, cpe_channels)
)
self.conv = spconv.SubMConv3d(
cpe_channels,
cpe_channels,
kernel_size=3,
bias=True,
indice_key=cpe_indice_key,
)
self.up = nn.Linear(cpe_channels, channels)
self.norm = norm_layer(channels)
def forward(self, point: Point):
point.feat = self.down(point.feat)
point.sparse_conv_feat = point.sparse_conv_feat.replace_feature(point.feat)
point.sparse_conv_feat = self.conv(point.sparse_conv_feat)
point.feat = point.sparse_conv_feat.features
point.feat = self.up(point.feat)
point.feat = self.norm(point.feat)
return point
class Block(PointModule):
def __init__(
self,
channels,
num_heads,
patch_size=48,
mlp_ratio=4.0,
qk_norm=False,
qkv_bias=True,
qk_scale=None,
attn_drop=0.0,
proj_drop=0.0,
drop_path=0.0,
layer_scale=None,
norm_layer=nn.LayerNorm,
act_layer=nn.GELU,
pre_norm=True,
order_index=0,
cpe_indice_key=None,
enable_rpe=False,
enable_flash=True,
upcast_attention=True,
upcast_softmax=True,
enable_cpe=True,
cpe_channels=96,
rope_base=10,
rope_jitter=None,
rope_rescale=None,
):
super().__init__()
self.channels = channels
self.pre_norm = pre_norm
if enable_cpe:
self.cpe = BottleneckCPE(
channels=channels,
cpe_channels=cpe_channels,
cpe_indice_key=cpe_indice_key,
norm_layer=norm_layer,
)
else:
self.cpe = None
self.norm1 = PointSequential(norm_layer(channels))
self.ls1 = PointSequential(
LayerScale(channels, init_values=layer_scale)
if layer_scale is not None
else nn.Identity()
)
self.attn = SerializedAttention(
channels=channels,
patch_size=patch_size,
num_heads=num_heads,
qk_norm=qk_norm,
qkv_bias=qkv_bias,
qk_scale=qk_scale,
attn_drop=attn_drop,
proj_drop=proj_drop,
order_index=order_index,
enable_rpe=enable_rpe,
enable_flash=enable_flash,
upcast_attention=upcast_attention,
upcast_softmax=upcast_softmax,
rope_base=rope_base,
rope_jitter=rope_jitter,
rope_rescale=rope_rescale,
)
self.norm2 = PointSequential(norm_layer(channels))
self.ls2 = PointSequential(
LayerScale(channels, init_values=layer_scale)
if layer_scale is not None
else nn.Identity()
)
self.mlp = PointSequential(
MLP(
in_channels=channels,
hidden_channels=int(channels * mlp_ratio),
out_channels=channels,
act_layer=act_layer,
drop=proj_drop,
)
)
self.drop_path = PointSequential(
DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
)
def forward(self, point: Point):
if self.cpe is not None:
with torch.autocast(device_type="cuda", enabled=False):
orig_dtype = point.feat.dtype
# => float32
point.feat = point.feat.to(dtype=torch.float32)
point.sparse_conv_feat = point.sparse_conv_feat.replace_feature(
point.sparse_conv_feat.features.to(torch.float32)
)
# perform fwd
shortcut = point.feat
point = self.cpe(point)
point.feat = shortcut + point.feat
# => original dtype
point.feat = point.feat.to(orig_dtype)
shortcut = point.feat
if self.pre_norm:
point = self.norm1(point)
point = self.drop_path(self.ls1(self.attn(point)))
point.feat = shortcut + point.feat
if not self.pre_norm:
point = self.norm1(point)
shortcut = point.feat
if self.pre_norm:
point = self.norm2(point)
point = self.drop_path(self.ls2(self.mlp(point)))
point.feat = shortcut + point.feat
if not self.pre_norm:
point = self.norm2(point)
point.sparse_conv_feat = point.sparse_conv_feat.replace_feature(point.feat)
return point
class GridPooling(PointModule):
def __init__(
self,
in_channels,
out_channels,
stride=2,
norm_layer=None,
act_layer=None,
reduce="max",
shuffle_orders=True,
traceable=True, # record parent and cluster
):
super().__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.stride = stride
assert reduce in ["sum", "mean", "min", "max"]
self.reduce = reduce
self.shuffle_orders = shuffle_orders
self.traceable = traceable
self.proj = nn.Linear(in_channels, out_channels)
if norm_layer is not None:
self.norm = PointSequential(norm_layer(out_channels))
if act_layer is not None:
self.act = PointSequential(act_layer())
def forward(self, point: Point):
if "grid_coord" in point.keys():
grid_coord = point.grid_coord
elif {"coord", "grid_size"}.issubset(point.keys()):
grid_coord = torch.div(
point.coord - point.coord.min(0)[0],
point.grid_size,
rounding_mode="trunc",
).int()
else:
raise AssertionError(
"[gird_coord] or [coord, grid_size] should be include in the Point"
)
grid_coord = torch.div(grid_coord, self.stride, rounding_mode="trunc")
grid_coord = grid_coord | point.batch.view(-1, 1) << 48
grid_coord, cluster, counts = torch.unique(
grid_coord,
sorted=True,
return_inverse=True,
return_counts=True,
dim=0,
)
grid_coord = grid_coord & ((1 << 48) - 1)
# indices of point sorted by cluster, for torch_scatter.segment_csr
_, indices = torch.sort(cluster)
# index pointer for sorted point, for torch_scatter.segment_csr
idx_ptr = torch.cat([counts.new_zeros(1), torch.cumsum(counts, dim=0)])
# head_indices of each cluster, for reduce attr e.g. code, batch
head_indices = indices[idx_ptr[:-1]]
point_dict = Dict(
feat=torch_scatter.segment_csr(
self.proj(point.feat)[indices], idx_ptr, reduce=self.reduce
),
coord=torch_scatter.segment_csr(
point.coord[indices], idx_ptr, reduce="mean"
),
grid_coord=grid_coord,
batch=point.batch[head_indices],
)
if "origin_coord" in point.keys():
point_dict["origin_coord"] = torch_scatter.segment_csr(
point.origin_coord[indices], idx_ptr, reduce="mean"
)
if "condition" in point.keys():
point_dict["condition"] = point.condition
if "context" in point.keys():
point_dict["context"] = point.context
if "name" in point.keys():
point_dict["name"] = point.name
if "split" in point.keys():
point_dict["split"] = point.split
if "color" in point.keys():
point_dict["color"] = torch_scatter.segment_csr(
point.color[indices], idx_ptr, reduce="mean"
)
if "segment_motif" in point.keys():
point_dict["segment_motif"] = point.segment_motif[head_indices]
if "grid_size" in point.keys():
point_dict["grid_size"] = point.grid_size * self.stride
if self.traceable:
point_dict["pooling_inverse"] = cluster
point_dict["pooling_parent"] = point
order = point.order
point = Point(point_dict)
if self.norm is not None:
point = self.norm(point)
if self.act is not None:
point = self.act(point)
point.serialization(order=order, shuffle_orders=self.shuffle_orders)
point.sparsify()
return point
class GridUnpooling(PointModule):
def __init__(
self,
in_channels,
skip_channels,
out_channels,
norm_layer=None,
act_layer=None,
traceable=False, # record parent and cluster
):
super().__init__()
self.proj = PointSequential(nn.Linear(in_channels, out_channels))
self.proj_skip = PointSequential(nn.Linear(skip_channels, out_channels))
if norm_layer is not None:
self.proj.add(norm_layer(out_channels))
self.proj_skip.add(norm_layer(out_channels))
if act_layer is not None:
self.proj.add(act_layer())
self.proj_skip.add(act_layer())
self.traceable = traceable
def forward(self, point):
assert "pooling_parent" in point.keys()
assert "pooling_inverse" in point.keys()
parent = point.pop("pooling_parent")
inverse = point.pooling_inverse
feat = point.feat
parent = self.proj_skip(parent)
parent.feat = parent.feat + self.proj(point).feat[inverse]
parent.sparse_conv_feat = parent.sparse_conv_feat.replace_feature(parent.feat)
if self.traceable:
point.feat = feat
parent["unpooling_parent"] = point
parent["unpooling_inverse"] = inverse
return parent
class Embedding(PointModule):
def __init__(
self,
in_channels,
embed_channels,
norm_layer=None,
act_layer=None,
mask_token=False,
):
super().__init__()
self.in_channels = in_channels
self.embed_channels = embed_channels
self.stem = PointSequential(linear=nn.Linear(in_channels, embed_channels))
if norm_layer is not None:
self.stem.add(norm_layer(embed_channels), name="norm")
if act_layer is not None:
self.stem.add(act_layer(), name="act")
if mask_token:
self.mask_token = nn.Parameter(torch.zeros(1, embed_channels))
else:
self.mask_token = None
def forward(self, point: Point):
point = self.stem(point)
if "mask" in point.keys():
point.feat = torch.where(
point.mask.unsqueeze(-1),
self.mask_token.to(point.feat.dtype),
point.feat,
)
return point
