aidsorb.litmodules

LightningModule’s for use with PyTorch Lightning.

class aidsorb.litmodules.PCDLit(model, criterion, metric, config_optimizer=None, config_scheduler=None)[source]

Bases: LightningModule

LightningModule for supervised learning on point clouds.

Note

  • *_step methods expect a batch of the form (x, y), where:

  • criterion must have signature criterion(input=preds, target=y).

  • Dictionaries passed as arguments are not deep copied. To avoid side effects, consider passing a deep copy.

Tip

You can use 'val_<MetricName>' as the quantity to monitor. For example, if:

from torchmetrics import R2Score, MeanAbsoluteError, MetricCollection
metric = MetricCollection(R2Score(), MeanAbsoluteError())

and you want to monitor R2Score, configure the ModelCheckpoint as following:

from lightning.pytorch.callbacks import ModelCheckpoint
checkpoint_callback = ModelCheckpoint(monitor='val_R2Score', mode='max', ...)
Parameters:

  • model (torch.nn.Module) – Architecture for point cloud processing.

  • criterion (callable) – Loss function to be optimized during training.

  • metric (torchmetrics.MetricCollection) – Metric(s) to be logged and optionally monitored. All metric(s) are logged based on Lightning’s default logging behavior. For more information, see logging.

  • config_optimizer (dict, default=None) –

    Dictionary for configuring optimizer. If None, the Adam optimizer with default hyperparameters is used.

    • 'name' optimizer’s class name str

    • 'hparams' optimizer’s hyperparameters dict

  • config_scheduler (dict, optional) –

    Dictionary for configuring learning rate scheduler.

Examples

>>> import torch
>>> from aidsorb.modules import PointNetClsHead, PointNet
>>> from torchmetrics import MetricCollection, R2Score, MeanAbsoluteError as MAE

>>> model = PointNet(head=PointNetClsHead(n_outputs=10))
>>> criterion, metric = torch.nn.MSELoss(), MetricCollection(R2Score(), MAE())

>>> # Adam optimizer with default hyperparameters, no scheduler.
>>> litmodel = PCDLit(model, criterion, metric)

>>> # Custom optimizer and scheduler.
>>> config_optimizer = {
... 'name': 'SGD',
... 'hparams': {'lr': 0.1},
... }
>>> config_scheduler = {
... 'name': 'StepLR',
... 'hparams': {'step_size': 2},
... 'config': {'interval': 'step'},
... }
>>> litmodel = PCDLit(model, criterion, metric, config_optimizer, config_scheduler)

>>> # Forward pass.
>>> x = torch.randn(32, 5, 100)
>>> litmodel(x).shape
torch.Size([32, 10])
configure_optimizers()[source]

Configure optimizer and optionally learning rate scheduler.

Warning

Parameters for which requires_grad=False are excluded from optimization.

Returns:

optimizers – Single optimizer if scheduler=None, else dictionary with keys: 'optimizer' and 'lr_scheduler'.

Return type:

Optimizer or dict

Examples

>>> import torch
>>> from torchmetrics import MetricCollection, R2Score
>>> criterion = torch.nn.MSELoss()
>>> metric = MetricCollection(R2Score())

>>> model = torch.nn.Linear(2, 2)
>>> litmodel = PCDLit(model, criterion, metric)
>>> litmodel.configure_optimizers()
...
Adam (
    ...
)

>>> model = torch.nn.Linear(4, 4)
>>> _ = model.requires_grad_(False)
>>> litmodel = PCDLit(model, criterion, metric)
>>> litmodel.configure_optimizers()
Traceback (most recent call last):
    ...
ValueError: optimizer got an empty parameter list
forward(x)[source]

Run forward pass (forward method) of model.

Parameters:

x (Any)

Return type:

Any

predict_step(batch, batch_idx)[source]

Return predictions on a single batch.

Parameters:
Return type:

Any

test_step(batch, batch_idx)[source]

Make predictions on a single batch from the test set and log metric(s).

Parameters:
Return type:

None

training_step(batch, batch_idx)[source]

Return loss on a single batch from the train set and log metric(s).

Parameters:
Return type:

Tensor

validation_step(batch, batch_idx)[source]

Make predictions on a single batch from the validation set and log metric(s).

Parameters:
Return type:

None