module

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: 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: _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

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: 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:

• step –

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

• 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(batch: Dict[str, Any], stage: str) -> Dict[str, Any]

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