datamodule

Bases: LightningDataModule

LightningDataModule wrapping point cloud datasets with the packed-batch collate.

Each dataset is passed through as-is. To lengthen an epoch, wrap the training dataset with torch_pointcloud.datasets.RepeatDataset(dataset, k=...) before passing it in.

Loaders are built with torch_pointcloud.utils.data.PointCloudDataLoader, which collates to the packed-batch torch_pointcloud.utils.data.collate. Collation specs are never read off the dataset (transforms rewrite keys downstream); pass stack_keys / cat_keys to control how per-scene ground truth is batched (e.g. cat_keys=("box",) for a detection dataset). shuffle is forced to True for train and False for val/test. Without a val_dataset, val_dataloader returns an empty list so Trainer.fit runs train-only.

Parameters:

• train_dataset (Optional[Dataset], default: None ) –

  Dataset for the training loop.

• train_ratios (Optional[Sequence[int]], default: None ) –

  One positive integer sampling weight per child dataset of a ConcatDataset train set. When set, the train loader draws single-dataset batches interleaved by these ratios via torch_pointcloud.datasets.SingleDatasetBatchSampler, so every batch stays single-domain (required by per-dataset normalization such as PDNorm). Leave None for a single dataset.

• val_dataset (Optional[Dataset], default: None ) –

  Dataset for the validation loop.

• test_dataset (Optional[Dataset], default: None ) –

  Dataset for the test loop.

• 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.

• batch_size (int, default: 1 ) –

  Number of point clouds per batch.

• eval_batch_size (Optional[int], default: None ) –

  Batch size of the val/test loaders; defaults to batch_size. Evaluation often runs full-resolution scenes while training runs crops, so the two memory envelopes differ (a benchmark protocol is typically one scene per batch).

• num_workers (int, default: 0 ) –

  Number of worker processes for data loading.

• pin_memory (bool, default: False ) –

  Pin tensors in pinned (page-locked) memory before transfer.

• drop_last (bool, default: True ) –

  Drop the last incomplete batch. Applied to the train loader only; val/test always set drop_last=False.

• timeout (float, default: 0.0 ) –

  Timeout for collecting a batch from the workers.

• prefetch_factor (Optional[int], default: None ) –

  Number of batches each worker prefetches ahead.

• persistent_workers (bool, default: False ) –

  Keep workers alive between epochs.

• pin_memory_device (str, default: '' ) –

  Target device for pin_memory (e.g. "cuda").

Methods:

• setup –

  Graft the model's evaluation transform onto datasets that have none, and collate its inverse key.

• configure_dataloader –

  Build a PointCloudDataLoader over the dataset with the module's collation and worker settings.

• train_dataloader –

  Build the shuffled train loader, drawing single-dataset batches when train_ratios is set.

• val_dataloader –

  Build the unshuffled validation loader, or an empty list when no validation set was given.

• test_dataloader –

  Build the unshuffled test loader, falling back to the validation set when no test set was given.

setup(stage: str) -> None

configure_dataloader(
    dataset: Optional[Dataset],
    *,
    shuffle: bool,
    drop_last: bool,
    batch_size: Optional[int] = None,
    sampler: Optional[Union[Sampler, Iterable]] = None,
    batch_sampler: Optional[
        Union[Sampler, Iterable]
    ] = None,
) -> DataLoader

train_dataloader() -> DataLoader

val_dataloader() -> Union[DataLoader, List[DataLoader]]

test_dataloader() -> DataLoader