data

Data loading: standard sample keys, packed-batch collation and row offsets, worker seeding, and the data loader.

Classes:

DataKeys –

Standard keys of a sample dict, shared by the datasets, transforms and models.
PointCloudDataLoader –

DataLoader that batches point clouds with the packed-batch collate by default.

Functions:

collate –

Collate a list of point-cloud sample dicts into one batched dict.
mix_pairs –

Merge consecutive pairs of samples with a pairwise mix, with probability p per batch.
select_inputs –

Gather the positional inputs of a model's forward from a batch dict, in keys order.
set_random_states –

Re-seed every seeded random object reachable from obj, each with its own seed.
seed_worker –

worker_init_fn re-seeding the dataset's seeded random objects from the worker's seed.
offset_index –

Offset per-element row indices into the packed row layout of a collated batch.

DataKeys [source] ¶

Bases: StrEnum

Standard keys of a sample dict, shared by the datasets, transforms and models.

Members are plain strings, so a key is usable wherever a literal is (e.g. data[DataKeys.POS] and data["pos"] address the same entry). Keys are singular (box, label), while the metric records (Boxes3D, Detection3D) use plural fields (boxes, labels).

PointCloudDataLoader [source] ¶

PointCloudDataLoader(
    dataset: Dataset,
    *,
    batch_from: str = POS,
    batch_key: str = BATCH,
    stack_keys: Optional[Sequence[str]] = None,
    cat_keys: Optional[Sequence[str]] = None,
    mix: Optional[
        Callable[
            [Dict[str, Any], Dict[str, Any]], Dict[str, Any]
        ]
    ] = None,
    mix_prob: float = 0.0,
    **kwargs: Any,
)

Bases: DataLoader

DataLoader that batches point clouds with the packed-batch collate by default.

Wraps torch.utils.data.DataLoader, defaulting collate_fn to collate. How keys collate is set by the spec arguments (batch_from / batch_key for the per-point index, stack_keys for dense per-scene ground truth, cat_keys for ragged per-scene ground truth). These are supplied by the caller, never read off the dataset: transforms rewrite the key set downstream of the dataset (a box key may be derived from an object by a transform), so only the code building the loader knows which keys must stack or cat. Passing collate_fn=... via the usual DataLoader kwarg overrides the spec.

Workers are re-seeded with seed_worker before any worker_init_fn passed in kwargs, so seeded random transforms draw different numbers in every worker and epoch.

Parameters:

dataset (Dataset) –

The dataset to load from.
batch_from (str, default: POS ) –

Key whose leading dimension defines the per-point batch index.
batch_key (str, default: BATCH ) –

Output key for the per-point batch index.
stack_keys (Optional[Sequence[str]], default: None ) –

Keys collated by stacking to a leading batch dim instead of concatenating.
cat_keys (Optional[Sequence[str]], default: None ) –

Packed keys that additionally emit a batch_<key> per-element scene index.
mix (Optional[Callable[[Dict[str, Any], Dict[str, Any]], Dict[str, Any]]], default: None ) –

Optional pairwise mix (e.g. Mix3D) merging consecutive samples of a batch before collation, with probability mix_prob per batch; see mix_pairs.
mix_prob (float, default: 0.0 ) –

Probability of mixing a batch.
**kwargs (Any, default: {} ) –

Forwarded to torch.utils.data.DataLoader (batch_size, shuffle, collate_fn, ...).

collate [source] ¶

collate(
    data_list: List[Dict[str, Any]],
    batch_from: str = POS,
    batch_key: str = BATCH,
    stack_keys: Optional[KeyCollection] = None,
    cat_keys: Optional[KeyCollection] = None,
) -> Dict[str, Any]

Collate a list of point-cloud sample dicts into one batched dict.

By default every per-point tensor is concatenated PyG-style along dim 0 (packed), scalars are stacked, and a per-point batch_key index is synthesized from batch_from. Two extra knobs say how specific keys collate instead:

stack_keys: stack to a new leading batch dim (\((M, \cdot) \to (B, M, \cdot)\), \((N, \cdot) \to (B, N, \cdot)\)) rather than concatenating. Used for fixed-size per-scene tensors that must keep their scene dimension (a plain cat would flatten them).
cat_keys: keep these packed (cat) but additionally emit a batch_<key> scene index mirroring batch_key. Used for ragged per-scene ground truth such as box \((K, 8)\) -> batch_box \((K,)\).

Every key must be present in every sample; a key missing from a sample raises a ValueError. stack_keys / cat_keys entries absent from all samples are ignored. A key may appear in only one of stack_keys / cat_keys; overlapping entries raise a ValueError.

Parameters:

data_list (List[Dict[str, Any]]) –

List of sample dicts.
batch_from (str, default: POS ) –

Key whose leading dimension defines the per-point batch index.
batch_key (str, default: BATCH ) –

Output key for the per-point batch index.
stack_keys (Optional[KeyCollection], default: None ) –

Keys collated by stacking to a leading batch dim instead of concatenating.
cat_keys (Optional[KeyCollection], default: None ) –

Packed keys that additionally emit a batch_<key> per-element scene index.

Returns:

Dict[str, Any] –

A single batched dict.

mix_pairs [source] ¶

mix_pairs(
    data_list: List[Dict[str, Any]],
    mix: Callable[
        [Dict[str, Any], Dict[str, Any]], Dict[str, Any]
    ],
    p: float = 1.0,
    generator: Optional[Generator] = None,
) -> List[Dict[str, Any]]

Merge consecutive pairs of samples with a pairwise mix, with probability p per batch.

When the draw succeeds, samples \((0, 1), (2, 3), \ldots\) of the batch are each merged into one scene by mix (Mix3D concatenates them), so the batch holds half as many scenes with all their points; an odd last sample passes through. Otherwise the batch is returned unchanged.

Parameters:

data_list (List[Dict[str, Any]]) –

The samples of a batch, before collation.
mix (Callable[[Dict[str, Any], Dict[str, Any]], Dict[str, Any]]) –

Pairwise transform called as mix(first, second).
p (float, default: 1.0 ) –

Probability of mixing the batch.
generator (Optional[Generator], default: None ) –

Random generator of the draw; None uses the global one.

Returns:

List[Dict[str, Any]] –

The (possibly mixed) samples.

Example

>>> import torch
>>> import torch_pointcloud.transforms as T
>>> scenes = [{"pos": torch.rand(10 * (i + 1), 3)} for i in range(4)]
>>> mixed = mix_pairs(scenes, T.Mix3D(keys="pos", instance_key=None, p=1.0), p=1.0)
>>> [scene["pos"].shape[0] for scene in mixed]
[30, 70]

select_inputs [source] ¶

select_inputs(
    data: Dict[str, Any], keys: Sequence[str]
) -> List[Any]

Gather the positional inputs of a model's forward from a batch dict, in keys order.

A dotted key reads an attribute of a batch value (e.g. octree.depth). A missing x resolves to None, since models accept x=None and build their features from the positions; any other missing key raises.

Parameters:

data (Dict[str, Any]) –

Collated batch dict, e.g. from collate or PointCloudDataLoader.
keys (Sequence[str]) –

Batch keys forward takes positionally, typically info["input_keys"] from create_model(..., return_info=True).

Returns:

List[Any] –

The values to unpack into the model's forward.

Raises:

KeyError –

If a key other than x is missing from data.

Example

import torch

import torch_pointcloud as tp
from torch_pointcloud.utils.data import collate, select_inputs

model, info = tp.create_model("pointnet.modelnet40", return_info=True)
samples = [info["transform"]({"pos": torch.rand(2048, 3), "normal": torch.rand(2048, 3)}) for _ in range(2)]
batch = collate(samples)
logits = model(*select_inputs(batch, info["input_keys"]))

set_random_states [source] ¶

set_random_states(obj: Any, seed: int) -> int

Re-seed every seeded random object reachable from obj, each with its own seed.

A random object exposes set_random_state(seed=...) and its generator as R (e.g. a Randomizable transform). Objects are found through sequences, mappings and instance attributes, such as a dataset's transform. Objects drawing from the global generator (R is None) are left as they are.

Parameters:

obj (Any) –

Root object, typically a dataset.
seed (int) –

Seed of the first object; the next ones get seed + 1, seed + 2, ...

Returns:

int –

The seed following the last one used.

Example

from torch_pointcloud.transforms import RandomJitter
from torch_pointcloud.utils.data import set_random_states

jitter = RandomJitter(keys="pos", seed=0)
set_random_states({"transform": jitter}, seed=7)

seed_worker [source] ¶

seed_worker(worker_id: int) -> None

worker_init_fn re-seeding the dataset's seeded random objects from the worker's seed.

Every worker receives a copy of the dataset, so a seeded transform would replay the same numbers in every worker. This re-seeds them from torch.utils.data.get_worker_info().seed, which differs per worker and per epoch and is reproducible under torch.manual_seed. PointCloudDataLoader applies it by default; pass it as worker_init_fn to a plain DataLoader.

Parameters:

worker_id (int) –

Index of the worker, as passed by the DataLoader.

Example

import torch
from torch.utils.data import DataLoader, TensorDataset

from torch_pointcloud.utils.data import seed_worker

loader = DataLoader(TensorDataset(torch.arange(8)), num_workers=2, worker_init_fn=seed_worker)
batch = next(iter(loader))

offset_index [source] ¶

offset_index(
    index: Tensor, index_batch: Tensor, batch: Tensor
) -> Tensor

Offset per-element row indices into the packed row layout of a collated batch.

collate concatenates a row map such as inverse or index as-is, so the entries of each batch element still address that element's own rows. This shifts every entry by the number of rows of the elements collated before it.

Parameters:

index (Tensor) –

Per-element row indices.
index_batch (Tensor) –

Batch index of each entry of index, the batch_<key> tensor collate emits for cat_keys.
batch (Tensor) –

Batch index of the rows index addresses (batch for an inverse map, batch_origin_pos for an index map).

Returns:

Tensor –

The offset row indices.

Shape

index: \((M,)\)
index_batch: \((M,)\)
batch: \((N,)\)
output: \((M,)\)

Example

import torch
from torch_pointcloud.utils.data import offset_index

inverse = torch.tensor([0, 1, 1, 0, 2, 2, 1])
batch_inverse = torch.tensor([0, 0, 0, 1, 1, 1, 1])
batch = torch.tensor([0, 0, 1, 1, 1])
offset_index(inverse, batch_inverse, batch)  # tensor([0, 1, 1, 2, 4, 4, 3])