vis4d.pl.training_module

LightningModule that wraps around the models, losses and optims.

Classes

TrainingModule(model_cfg, optimizers_cfg, ...)

LightningModule that wraps around the vis4d implementations.

class TrainingModule(model_cfg, optimizers_cfg, loss_module, train_data_connector, test_data_connector, hyper_parameters=None, seed=-1, ckpt_path=None, compute_flops=False)

LightningModule that wraps around the vis4d implementations.

This is a wrapper around the vis4d implementations that allows to use pytorch-lightning for training and testing.

Initialize the TrainingModule.

Parameters:

• model_cfg (ConfigDict) – The model config.

• optimizers_cfg (list[OptimizerConfig]) – The optimizers config.

• loss_module (None | LossModule) – The loss module.

• train_data_connector (None | DataConnector) – The data connector to use.

• test_data_connector (None | DataConnector) – The data connector to use.

• data_connector – The data connector to use.

• hyper_parameters (DictStrAny | None, optional) – The hyper parameters to use. Defaults to None.

• seed (int, optional) – The integer value seed for global random state. Defaults to -1. If -1, a random seed will be generated.

• ckpt_path (str, optional) – The path to the checkpoint to load. Defaults to None.

• compute_flops (bool, optional) – If to compute the FLOPs of the model. Defaults to False.

configure_optimizers()

Return the optimizer to use.

Return type:

Any

forward(data)

Forward pass through the model.

Return type:

Any

lr_scheduler_step(scheduler, metric=None)

Perform a step on the lr scheduler.

Return type:

None

setup(stage)

Setup the model.

Return type:

None

test_step(batch, batch_idx, dataloader_idx=0)

Perform a single test step.

Return type:

Dict[str, Any]

training_step(batch, batch_idx)

Perform a single training step.

Return type:

Any

validation_step(batch, batch_idx, dataloader_idx=0)

Perform a single validation step.

Return type:

Dict[str, Any]