discotime.models.models module

class discotime.models.models.LitSurvModule(config: ModelConfig)

Bases: LightningModule

configure_optimizers() → Optimizer

property data_cuts: Tensor

property datamodule: LitSurvDataModule

property eval_grid: Tensor

forward(x: Tensor) → Tensor

property model: Module

on_load_checkpoint(checkpoint: Dict[str, Any]) → None

on_save_checkpoint(checkpoint: Dict[str, Any]) → None

on_test_epoch_end() → None

on_validation_epoch_end() → None

setup(stage)

Called at the beginning of fit (train + validate), validate, test, or predict. This is where the PyTorch model gets instantiated.

Parameters:

stage (str) – either "fit", "validate" "test" or "predict"

test_step(batch: tuple[torch.Tensor, ...], batch_idx: int) → Tensor

training_step(batch: tuple[torch.Tensor, ...], batch_idx: int) → Tensor

validation_step(batch: tuple[torch.Tensor, ...], batch_idx: int) → Tensor

class discotime.models.models.ModelConfig(*, learning_rate: float = 0.001, activation_function: str = 'SiLU', n_sequential_blocks: int = 2, n_hidden_units: int = 20, dropout_rate: float | None = None, batch_normalization: bool = True, use_skip_connections: bool = False, evaluation_grid_size: int = 50, evaluation_grid_cuts: int | None = None, n_time_bins: int | None = None, n_risks: int | None = None)

Bases: object

activation_function: str = 'SiLU'

Name of activation function (str). The name needs to match one of the activation functions in the torch.nn module.

batch_normalization: bool = True

Use batch normalization in each block? Default is True.

dropout_rate: float | None = None

Should nn.Dropout() be used in each block, and if so what is the rate of dropout? If None then dropout is not used. Default is None.

evaluation_grid_cuts: int | None = None

Either the length of (int) or the specific grid (seq of floats) at which model metrics are calculated. If an integer n is passed, then n evenly distributed timepoint are chosen from the range of the data.

evaluation_grid_size: int = 50

Either the length of (int) or the specific grid (seq of floats) at which model metrics are calculated. If an integer n is passed, then n evenly distributed timepoint are chosen from the range of the data.

learning_rate: float = 0.001

n_hidden_units: int = 20

Number of neurons in each of the hidden layers.

For ease of of use, the size of each hidden layer have been constrained to to be the exact same size as all the others. Default is 20.

n_risks: int | None = None

How many competing risks are we dealing with?

If None, which is the default, then we try to get the information from the attached datamodule during setup.

n_sequential_blocks: int = 2

Number of neural-network layer blocks.

Each block is modelled after a ResNet skip-block and contains the following key elements:

Sequential(
    (0): LazyBatchNorm1d()
    (1): LazyLinear()
    (2): DropOut()
    (3): SiLU()
    (4): LazyLinear()
    (5): DropOut()
)

Per default the output of the model is self.act_fn(x + self.net(x)) but the skip part can be removed by setting use_skip_connections to False. n_hidden_units controls the size of the linear layers in each block. See components.Block for more details on the actual implementation.

n_time_bins: int | None = None

Number of time bins in discretization grid.

If None, which is the default, then we try to get the information from the attached datamodule during setup.

use_skip_connections: bool = False

Toggle the use of skip-connections in the model blocks.