pdnorm

Prompt-driven normalization (PDNorm) routing each batch through a per-condition norm.

Classes:

PDNorm –

Prompt-driven normalization that routes each batch through a per-condition norm.

PDNorm ¶

PDNorm(
    channels: int,
    conditions: Sequence[str],
    norm: Union[str, Callable, None] = "batch_norm",
    decouple: bool = True,
    *,
    dim: int = 1,
    **norm_kwargs: Any,
)

Bases: Module

Prompt-driven normalization that routes each batch through a per-condition norm.

Multi-dataset joint training feeds batches drawn from a single dataset at a time, identified by a string condition. When decouple is True, PDNorm holds one independent inner norm per condition (an nn.ModuleList indexed by conditions.index(condition)), so each dataset keeps its own running statistics and affine parameters. When decouple is False, a single shared norm is applied regardless of the condition.

Inner norms are built with create_norm, so any name / class / instance accepted there is valid. The decoupled layout stores children under norm.{i} keys, matching the order of conditions.

Introduced in Towards Large-scale 3D Representation Learning with Multi-dataset Point Prompt Training

Parameters:

channels (int) –

Number of feature channels.
conditions (Sequence[str]) –

Ordered condition names; index \(i\) selects norm.{i} when decoupled.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Inner norm passed to create_norm (name, class, instance, or None).
decouple (bool, default: True ) –

If True, use one norm per condition; if False, share a single norm.
dim (int, default: 1 ) –

Spatial dimensionality hint forwarded to each inner create_norm (see its dim argument).
**norm_kwargs (Any, default: {} ) –

Extra keyword arguments forwarded to each inner norm constructor.

Shape

Input: \((N, C)\) packed features with \(C =\) channels.
Output: \((N, C)\), same shape as the input.

Example

import torch
from torch_pointcloud.layers import PDNorm

norm = PDNorm(64, conditions=["ScanNet", "S3DIS"], norm="batch_norm")
x = torch.randn(32, 64)
y = norm(x, condition="S3DIS")
print(y.shape)