Overview

Bases: Callback

Exponentially decay BatchNorm momentum over training epochs.

Reference implementation: facebookresearch/votenet (train.py).

At the start of each training epoch, every nn.BatchNorm* module in pl_module.model has its momentum set to

with \(m_0\) the initial momentum, \(\gamma\) the decay rate, \(s\) the decay step (epochs) and \(m_\text{clip}\) the floor (torch_pointcloud.optim.bn_momentum).

Parameters:

• bn_momentum_init (float, default: 0.5 ) –

  Initial BatchNorm momentum \(m_0\).

• bn_decay_rate (float, default: 0.5 ) –

  Per-step multiplicative decay \(\gamma\).

• bn_decay_step (int, default: 20 ) –

  Number of epochs between decay steps \(s\).

• bn_momentum_clip (float, default: 0.001 ) –

  Lower bound on the momentum \(m_\text{clip}\).

Methods:

• on_train_epoch_start –

  Set the decayed momentum on every BatchNorm module of the model.

on_train_epoch_start(
    trainer: Trainer, pl_module: LightningModule
) -> None

Bases: Callback

Accumulate and log a torchmetrics metric over a validation/test epoch.

Model- and task-agnostic: the LightningModule's validation_step / test_step returns a {preds_key: ..., target_key: ...} dict (the repo's Lit* modules do), and this callback updates a torchmetrics Metric with it each batch, logging the epoch value as {stage}/{name}. List one per metric to plug accuracy, mIoU, mAP, ... onto any model.

metric is a ready torchmetrics Metric whose num_classes and ignore_index the caller sets to match the model and the criterion, so the callback holds it as-is and stays model-agnostic.

A step output may carry extra entries beyond preds_key / target_key (the Lit* modules' metric_input_keys passthrough); entries matching a parameter name of the metric's update signature (inspected once at construction) are forwarded as keyword arguments and the rest are ignored, so each metric declares the inputs it consumes.

A metric whose compute returns a dict (e.g. BoxMeanAveragePrecision returning mAP@0.25 / mAP@0.5) is logged one entry per key as {stage}/{key}; a scalar metric is logged as {stage}/{name}.

Parameters:

• metric (Metric) –

  A torchmetrics Metric, e.g. JaccardIndex(task="multiclass", num_classes=..., ignore_index=...).

• name (str) –

  Metric name; logged as {stage}/{name} (ignored for dict-valued metrics, whose keys name themselves).

• stages (Sequence[str], default: ('val', 'test') ) –

  Stages to score; each listed stage's *_step must return preds_key / target_key.

• preds_key (str, default: 'preds' ) –

  Key in the step output holding predictions (logits, probabilities, labels, or detections).

• target_key (str, default: 'target' ) –

  Key in the step output holding the targets.

• batch_key (Optional[str], default: None ) –

  Optional key in the step output holding the per-point shape index, forwarded as a third positional argument to metric.update (packed multi-shape metrics like InstancePartMeanIntersectionOverUnion need it). Leave None for two-argument metrics.

• prog_bar (bool, default: True ) –

  Whether to show the metric on the progress bar.

Methods:

• on_validation_epoch_start –

  Reset the metric when val is scored.

• on_validation_batch_end –

  Update the metric with the step output when val is scored.

• on_validation_epoch_end –

  Compute and log the metric when val is scored.

• on_test_epoch_start –

  Reset the metric when test is scored.

• on_test_batch_end –

  Update the metric with the step output when test is scored.

• on_test_epoch_end –

  Compute and log the metric when test is scored.

on_validation_epoch_start(
    trainer: Trainer, pl_module: LightningModule
) -> None

on_validation_batch_end(
    trainer: Trainer,
    pl_module: LightningModule,
    outputs: Any,
    batch: Dict[str, Any],
    batch_idx: int,
    dataloader_idx: int = 0,
) -> None

on_validation_epoch_end(
    trainer: Trainer, pl_module: LightningModule
) -> None

on_test_epoch_start(
    trainer: Trainer, pl_module: LightningModule
) -> None

on_test_batch_end(
    trainer: Trainer,
    pl_module: LightningModule,
    outputs: Any,
    batch: Dict[str, Any],
    batch_idx: int,
    dataloader_idx: int = 0,
) -> None

on_test_epoch_end(
    trainer: Trainer, pl_module: LightningModule
) -> None

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

Bases: Metric

Packed 3D-detection per-class average precision as a torchmetrics metric.

A stateful wrapper of box_average_precision: each update reduces one batch's packed predictions and ground truth to its compact box_matches record, and compute returns one AP/<class> entry per class plus their mean as mAP, each class matched at its own IoU threshold (the KITTI / nuScenes convention, e.g. Car@0.7 and Pedestrian/Cyclist@0.5). Targets may carry an ignore_mask so predictions overlapping an ignore region are not counted as false positives. An ignore_mask passed to update is used as the predictions' ignore_mask entry, excluding the flagged predictions from scoring entirely (the KITTI min-height rule). The state is a per-process list of records (not gathered across processes), so run detection validation on a single device.

Parameters:

• iou_per_class (Mapping[int, float]) –

  Mapping of class index to the IoU threshold used to match its boxes. Keys are coerced to int (YAML / OmegaConf mappings may arrive with string keys).

• class_names (Optional[Sequence[str]], default: None ) –

  Optional class names used in the returned keys (defaults to the class index).

• interpolation (Interpolation, default: 'all' ) –

  AP interpolation: "all" integrates the full precision-recall curve; "r11" / "r40" sample the KITTI 11- / 40-point recall grids.

• kwargs (Any, default: {} ) –

  Forwarded to torchmetrics.Metric.

Methods:

• update –

  Reduce one batch's packed predictions and ground truth to its box_matches record and append it.

• compute –

  Score the accumulated batches and return one AP/<class> entry per class plus their mAP.

update(
    preds: Detection3D,
    target: Boxes3D,
    ignore_mask: OptTensor = None,
) -> None

compute() -> Dict[str, float]

Bases: Metric

Packed 3D-detection mean average precision as a torchmetrics metric.

A stateful wrapper of box_average_precision: each update reduces one batch's packed predictions and ground truth to its compact box_matches record, and compute returns {"mAP@t": ...} (averaged over the classes present in the targets) for each IoU threshold. An ignore_mask passed to update is used as the predictions' ignore_mask entry, excluding the flagged predictions from scoring entirely (the KITTI min-height rule). The state is a per-process list of records (not gathered across processes), so run detection validation on a single device.

Parameters:

• iou_thresholds (Sequence[float], default: (0.25, 0.5) ) –

  IoU thresholds at which mAP@t is reported.

• interpolation (Interpolation, default: 'all' ) –

  AP interpolation: "all" integrates the full precision-recall curve; "r11" / "r40" sample the KITTI 11- / 40-point recall grids.

• kwargs (Any, default: {} ) –

  Forwarded to torchmetrics.Metric.

Methods:

• update –

  Reduce one batch's packed predictions and ground truth to its box_matches record and append it.

• compute –

  Score the accumulated batches and return one mAP@t entry per IoU threshold.

update(
    preds: Detection3D,
    target: Boxes3D,
    ignore_mask: OptTensor = None,
) -> None

compute() -> Dict[str, float]

Bases: Metric

Point-mask instance-segmentation AP as a torchmetrics metric.

A stateful wrapper of instance_average_precision: each update appends one scene's instance_matches record (the compact per-scene reduction of predicted masks against ground-truth instances), and compute returns {"AP/<class>": ..., "mAP": ..., "mAP@0.5": ..., "mAP@0.25": ...} following the standard indoor instance-segmentation protocol. The state is a per-process list of records (not gathered across processes), so run instance validation on a single device.

Parameters:

• num_classes (int) –

  Number of instance classes.

• class_names (Optional[Sequence[str]], default: None ) –

  Optional names for the AP/<class> keys; falls back to the class index.

• min_points (int, default: 100 ) –

  Minimum point count for a prediction or ground-truth instance to be scored; smaller ground-truth instances count as ignore regions.

• kwargs (Any, default: {} ) –

  Forwarded to torchmetrics.Metric.

Methods:

• update –

  Append one scene's instance_matches record.

• compute –

  Score the accumulated scene records and return the per-class average precisions and their mean.

update(match: InstanceMatches) -> None

compute() -> Dict[str, float]

Bases: Metric

ShapeNetPart instance / class mean IoU as a torchmetrics metric.

A stateful wrapper of part_intersection_over_union: each shape is scored only over the part labels its category owns (a part absent from both the prediction and the target counts as IoU \(1\)), per-category IoU sums and shape counts accumulate across update calls (summed across processes), and compute returns the protocol's two numbers: ins_mIoU (mean over shapes) and cls_mIoU (mean per category, then over the categories seen). The category of each shape is read off its target labels, since every category owns a disjoint part range.

Parameters:

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

  Part labels owned by each category; defaults to the 16-category / 50-part ShapeNetPart table (ShapeNetPart.seg_ids).

• restrict_to_category (bool, default: False ) –

  If True, the argmax of 2-D preds is taken over the shape's own category parts only (logits of the other parts are masked out), the protocol of PointNet, Point-MAE and Point-M2AE. The default is the global argmax over all parts (DGCNN, PointNeXt).

• kwargs (Any, default: {} ) –

  Forwarded to torchmetrics.Metric.

Methods:

• update –

  Score one packed batch of shapes.

• compute –

  Reduce the accumulated per-category IoU sums into ins_mIoU and cls_mIoU.

update(
    preds: Tensor, target: Tensor, batch: Tensor
) -> None

compute() -> Dict[str, Tensor]

Bases: Metric

The official nuScenes detection metrics as a torchmetrics metric.

A stateful wrapper of nuscenes_detection_metrics: each update appends one batch's packed predictions and ground truth together with the optional velocity, attribute and point-count extras, and compute returns the functional's flat dict (AP/<class>, mAP, the five TP errors and the NDS). The predictions' velocity entry and the velocity argument are appended to the boxes as \((v_x, v_y)\) columns, the \((M, 9)\) layout the functional scores; prediction attributes are derived from the accumulated prediction velocities in compute with the standard speed heuristic (nuscenes_velocity_attributes). Sample indices are offset per update by the number of samples seen so far (read off the batch tensors), so scenes of different updates never collide. The state is a per-process list of batches (not gathered across processes), so run detection validation on a single device.

Parameters:

• class_names (Sequence[str]) –

  Class name per label index; barrier and traffic_cone get their official special handling by name.

• class_ranges (Optional[Mapping[str, float]], default: None ) –

  Maximum BEV evaluation range per class name; defaults to the official ranges.

• dist_thresholds (Sequence[float], default: (0.5, 1.0, 2.0, 4.0) ) –

  Matching thresholds in meters the AP is averaged over.

• tp_threshold (float, default: 2.0 ) –

  Matching threshold in meters of the TP-error metrics.

• max_boxes_per_sample (int, default: 500 ) –

  Per-sample cap on scored predictions (highest scores kept).

• min_recall (float, default: 0.1 ) –

  Recall up to which the AP and TP-error curves are clipped.

• min_precision (float, default: 0.1 ) –

  Precision subtracted before the AP mean.

• kwargs (Any, default: {} ) –

  Forwarded to torchmetrics.Metric.

Methods:

• update –

  Append one batch's packed predictions and ground truth with the optional nuScenes extras.

• compute –

  Score the accumulated samples and return the official nuScenes metrics, NDS included.

update(
    preds: Detection3D,
    target: Boxes3D,
    *,
    velocity: OptTensor = None,
    num_points: OptTensor = None,
    attribute: OptTensor = None,
) -> None

compute() -> Dict[str, float]

Bases: LitModel

LightningModule for a point cloud classification model, built from the registry.

Parameters:

• name (str) –

  Registered classification model name (e.g. pointnet2-ssg.modelnet40.xu-yan); built via create_model(name, task="classification").

• **kwargs (Any, default: {} ) –

  Forwarded to the base LitModel (e.g. optimizer, scheduler, criterion) and create_model (e.g. pretrained=True, or registry-hparam overrides).

Methods:

• predict –

  Forward the batch and return the logits tensor (the predictor handed to the inferer).

• step –

  Predict the batch, log the criterion loss, and return the logits, targets and loss.

• training_step –

  Run the shared step and return the loss to optimize.

• validation_step –

  Return the evaluation predictions and targets, plus the batch's metric_input_keys.

• test_step –

  Return the inferer's predictions and targets, plus the batch's metric_input_keys.

• configure_optimizers –

  Build the configured optimizer (over param_groups when given) and its optional scheduler.

predict(batch: Dict[str, Any]) -> Tensor

step(
    batch: Dict[str, Any], stage: str
) -> Dict[str, Tensor]

training_step(
    batch: Dict[str, Any], batch_idx: int
) -> Tensor

validation_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

test_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

configure_optimizers() -> Any

Bases: LitModel

LightningModule for a 3D object detection model, built from the registry.

Detection breaks the shared classification/segmentation loop in three places, so this subclass overrides them and reuses the base for everything else (model construction, forward, optimizer wiring, training_step):

• the loss is either a ready-built nn.Module (the general case: an anchor / center / set-matching loss whose geometry params are set explicitly in config) or a factory completed at build time with the model's head-geometry params (the VoteNetLoss carve-out, which reads num_heading_bins, num_size_clusters, num_classes, mean_sizes off the model); without a criterion the module is evaluation-only (no loss is logged, training raises);
• step (training only) feeds the whole forward output and the batch to the loss, which returns a dict of named components (each logged), and reports the total loss;
• the eval steps are metric-driven, not loss-driven: they run the model's raw decode, postprocess it (optional min_points filter, drop boxes below score_threshold, per-class 3D nms3d at nms_iou on the rotated BEV IoU when nms_rotated, and the indoor per-class expansion when decode emits class_probs), and pair the result with the ground truth for a MetricCallback (e.g. BoxMeanAveragePrecision, BoxAveragePrecision); any other per-box decode entry (e.g. the nuScenes heads' velocity) is filtered alongside the boxes and kept in the predictions dict, and when the batch carries DataKeys.CALIB / DataKeys.IMAGE_SHAPE (stacked per-frame \((B, 3, 4)\) / \((B, 2)\)) the sub-25 px projected_ignore_mask of the surviving boxes is attached as the predictions' ignore_mask (the KITTI min-height rule); the ground-truth boxes are DataKeys.BOX packed as \((K, 7)\) rows \([c_x, c_y, c_z, d_x, d_y, d_z, \theta]\) (full extents, counter-clockwise heading \(\theta\)), with per-box classes under label_key. No loss is computed at validation, so a two-stage detector whose inference forward differs from its training forward (its eval output carries decoded proposals, not loss targets) validates cleanly.

A model is swappable as long as it returns a prediction dict from forward and a raw Detection3D from decode(output), and (for training) is paired with a criterion(output, data).

Parameters:

• name (str) –

  Registered detection model name; built via create_model(name, task="detection").

• criterion (Union[Module, Callable[..., Module], None], default: None ) –

  The training loss, in one of two forms. A ready-built nn.Module is used as-is (the general case: instantiate it in config with its geometry params, e.g. AnchorHeadLoss). A callable factory is completed with the model's head-geometry params, i.e. called as criterion(num_heading_bins=..., num_size_clusters=..., num_classes=..., mean_sizes=...) (the VoteNetLoss carve-out). Either way its forward(output, data) returns a dict whose loss entry is the total to optimize. Leave None to benchmark a detector whose training loss is not ported.

• score_threshold (float, default: 0.05 ) –

  Minimum score to keep a decoded box in the eval postprocess.

• nms_iou (float, default: 0.25 ) –

  IoU threshold of the per-class 3D NMS in the eval postprocess.

• nms_rotated (bool, default: False ) –

  Suppress on the exact rotated BEV IoU instead of the axis-aligned 3D IoU (see nms3d); the KITTI outdoor protocol.

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

  Optional minimum number of points inside a decoded box for it to be kept (the VoteNet / 3DETR indoor protocol uses \(5\)).

• label_key (str, default: LABEL ) –

  Batch-dict key of the per-box ground-truth class labels (defaults to DataKeys.LABEL).

• ignore_mask_key (Optional[str], default: None ) –

  Optional batch-dict key of a per-box ignore mask forwarded to the metric with the ground truth (KITTI-style ignore regions).

• **kwargs (Any, default: {} ) –

  Forwarded to create_model and the base (e.g. optimizer, scheduler, input_keys). The base's inferer does not apply: detection eval is the decode / postprocess path above.

Methods:

• predict –

  Forward the batch and return the logits tensor (the predictor handed to the inferer).

• training_step –

  Run the shared step and return the loss to optimize.

• validation_step –

  Return the evaluation predictions and targets, plus the batch's metric_input_keys.

• test_step –

  Return the inferer's predictions and targets, plus the batch's metric_input_keys.

• configure_optimizers –

  Build the configured optimizer (over param_groups when given) and its optional scheduler.

• step –

  Forward the batch, log every named loss component, and return the output and the total loss.

predict(batch: Dict[str, Any]) -> Tensor

training_step(
    batch: Dict[str, Any], batch_idx: int
) -> Tensor

validation_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

test_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

configure_optimizers() -> Any

step(batch: Dict[str, Any], stage: str) -> Dict[str, Any]

Bases: LightningModule

Shared base for the task-specific Lightning wrappers.

The task-specific subclasses build their model through create_model; this base only holds the built model (and its registered evaluation transform) and implements the shared train/val/test loop. input_keys, target_key and scheduler_interval are read from self.hparams, which the subclass populates via save_hyperparameters. Without an optimizer the module is evaluation-only (benchmark mode): Trainer.test works and Trainer.fit raises.

Parameters:

• name (str) –

  Registered model name; built via create_model(name, task=...).

• task (Task) –

  Registry of the model ("classification", "semantic-segmentation", "part-segmentation", "detection"); set by the subclass.

• optimizer (Optional[Callable[..., Optimizer]], default: None ) –

  A callable that takes parameters and returns an optimizer (a _partial_ target).

• scheduler (Optional[Callable[..., Any]], default: None ) –

  An optional callable that takes an optimizer and returns a learning-rate scheduler.

• criterion (Optional[Module], default: None ) –

  The loss module; defaults to CrossEntropyLoss.

• inferer (Optional[Inferer], default: None ) –

  Test-time inference strategy (e.g. TTAInferer, SlidingWindowInferer) run in place of the plain forward on test batches; defaults to SimpleInferer (one forward on the whole batch), so every test prediction goes through an inferer. Training and validation are unaffected. The inferer may return probabilities instead of logits (torchmetrics handles both), so no test/loss is logged.

• input_keys (Optional[Sequence[str]], default: None ) –

  Batch-dict keys passed positionally to the model's forward (see select_inputs); None (the default) uses the registered input_keys of the model.

• target_key (str, default: LABEL ) –

  Batch-dict key for the per-cloud label.

• metric_input_keys (Sequence[str], default: () ) –

  Batch-dict keys copied as-is into the validation_step / test_step output dict, alongside the predictions and targets, for metrics whose update consumes extra inputs (MetricCallback forwards the keys each metric declares). A listed key missing from the batch raises.

• scheduler_interval (str, default: 'epoch' ) –

  Whether the scheduler steps per "epoch" or "step".

• param_groups (Optional[Dict[str, Any]], default: None ) –

  Optional dict of kwargs forwarded to torch_pointcloud.optim.generate_param_groups.

• **kwargs (Any, default: {} ) –

  Forwarded to create_model (e.g. pretrained=True, or registry-hparam overrides).

Methods:

• predict –

  Forward the batch and return the logits tensor (the predictor handed to the inferer).

• step –

  Predict the batch, log the criterion loss, and return the logits, targets and loss.

• training_step –

  Run the shared step and return the loss to optimize.

• validation_step –

  Return the evaluation predictions and targets, plus the batch's metric_input_keys.

• test_step –

  Return the inferer's predictions and targets, plus the batch's metric_input_keys.

• configure_optimizers –

  Build the configured optimizer (over param_groups when given) and its optional scheduler.

predict(batch: Dict[str, Any]) -> Tensor

step(
    batch: Dict[str, Any], stage: str
) -> Dict[str, Tensor]

training_step(
    batch: Dict[str, Any], batch_idx: int
) -> Tensor

validation_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

test_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

configure_optimizers() -> Any

Bases: _LitSegmentationModel

LightningModule for a part segmentation model, built from the registry.

The model also reads the one-hot object category, which the registered input_keys of part segmentation models list (("x", "pos", "batch", "category")).

Parameters:

• name (str) –

  Registered part segmentation model name; built via create_model(name, task="part-segmentation").

• inverse_key (Optional[str], default: None ) –

  Batch-dict key of the source-to-predictor row map written by the transform (inverse; see the transforms module docs on sampling keys). When set, eval predictions are broadcast to source resolution (preds[batch[inverse_key]]) and scored against origin_target_key; the loss stays at predictor resolution against target_key. Multi-scene batches need the key in the loader's cat_keys so the per-scene maps can be offset into the packed layout; the datamodule adds it.

• origin_target_key (Optional[str], default: None ) –

  Batch-dict key of the source-resolution labels scored with inverse_key; required with it.

• target_key (str, default: SEGMENT ) –

  Batch-dict key of the per-point labels.

• **kwargs (Any, default: {} ) –

  Forwarded to the base LitModel (e.g. inferer, optimizer, criterion) and create_model (e.g. pretrained=True, or registry-hparam overrides).

Methods:

• predict –

  Forward the batch and return the logits tensor (the predictor handed to the inferer).

• step –

  Predict the batch, log the criterion loss, and return the logits, targets and loss.

• training_step –

  Run the shared step and return the loss to optimize.

• validation_step –

  Return the evaluation predictions and targets, plus the batch's metric_input_keys.

• test_step –

  Return the inferer's predictions and targets, plus the batch's metric_input_keys.

• configure_optimizers –

  Build the configured optimizer (over param_groups when given) and its optional scheduler.

predict(batch: Dict[str, Any]) -> Tensor

step(
    batch: Dict[str, Any], stage: str
) -> Dict[str, Tensor]

training_step(
    batch: Dict[str, Any], batch_idx: int
) -> Tensor

validation_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

test_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

configure_optimizers() -> Any

Bases: _LitSegmentationModel

LightningModule for a point cloud semantic segmentation model, built from the registry.

Parameters:

• name (str) –

  Registered semantic segmentation model name; built via create_model(name, task="semantic-segmentation").

• inverse_key (Optional[str], default: None ) –

  Batch-dict key of the source-to-predictor row map written by the transform (inverse; see the transforms module docs on sampling keys). When set, eval predictions are broadcast to source resolution (preds[batch[inverse_key]]) and scored against origin_target_key; the loss stays at predictor resolution against target_key. Multi-scene batches need the key in the loader's cat_keys so the per-scene maps can be offset into the packed layout; the datamodule adds it.

• origin_target_key (Optional[str], default: None ) –

  Batch-dict key of the source-resolution labels scored with inverse_key; required with it.

• target_key (str, default: SEGMENT ) –

  Batch-dict key of the per-point labels.

• **kwargs (Any, default: {} ) –

  Forwarded to the base LitModel (e.g. inferer, optimizer, criterion) and create_model (e.g. pretrained=True, or registry-hparam overrides).

Methods:

• predict –

  Forward the batch and return the logits tensor (the predictor handed to the inferer).

• step –

  Predict the batch, log the criterion loss, and return the logits, targets and loss.

• training_step –

  Run the shared step and return the loss to optimize.

• validation_step –

  Return the evaluation predictions and targets, plus the batch's metric_input_keys.

• test_step –

  Return the inferer's predictions and targets, plus the batch's metric_input_keys.

• configure_optimizers –

  Build the configured optimizer (over param_groups when given) and its optional scheduler.

predict(batch: Dict[str, Any]) -> Tensor

step(
    batch: Dict[str, Any], stage: str
) -> Dict[str, Tensor]

training_step(
    batch: Dict[str, Any], batch_idx: int
) -> Tensor

validation_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

test_step(
    batch: Dict[str, Any], batch_idx: int
) -> Dict[str, Any]

configure_optimizers() -> Any