Inference engines#

The module tsl.engines contains inference engines, i.e., modules meant to wrap deep models in order to ease training and inference phases. Every engine extends a LightningModule from Lightning.

Currently, there are two basic engines:

Predictor

Use this module when you are addressing the prediction task, i.e., inference of future observations.
Imputer

Use this module when you are addressing the imputation task, i.e., reconstruction of missing observations.

We suggest to extends this module and override some of the methods to meet your project-specific needs.

class Predictor(model: Optional[Module] = None, loss_fn: Optional[Callable] = None, scale_target: bool = False, metrics: Optional[Mapping[str, Metric]] = None, *, model_class: Optional[Type] = None, model_kwargs: Optional[Mapping] = None, optim_class: Optional[Type] = None, optim_kwargs: Optional[Mapping] = None, scheduler_class: Optional = None, scheduler_kwargs: Optional[Mapping] = None)[source]#

LightningModule to implement predictors.

Input data should follow the format [batch, steps, nodes, features].

Parameters:

model (torch.nn.Module, optional) – Model implementing the predictor. Ignored if argument model_class is not None. This argument should mainly be used for inference. (default: None)
model_class (type, optional) – Class of Module implementing the predictor. If not None, argument model will be ignored. (default: None)
model_kwargs (mapping, optional) – Dictionary of arguments to be forwarded to model_class at instantiation. (default: None)
optim_class (type, optional) – Class of Optimizer implementing the optimizer to be used for training the model. (default: None)
optim_kwargs (mapping, optional) – Dictionary of arguments to be forwarded to optim_class at instantiation. (default: None)
loss_fn (callable, optional) – Loss function to be used for training the model. (default: None)
scale_target (bool) – Whether to scale target before evaluating the loss. The metrics instead will always be evaluated in the original range. (default: False)
metrics (mapping, optional) – Set of metrics to be logged during train, val and test steps. The metric’s name will be automatically prefixed with the loop in which the metric is computed (e.g., metric mae will be logged as train_mae when evaluated during training). (default: None)
scheduler_class (type, optional) – Class of _LRScheduler implementing the learning rate scheduler to be used during training. (default: None)
scheduler_kwargs (mapping, optional) – Dictionary of arguments to be forwarded to scheduler_class at instantiation. (default: None)

load_model(filename: str)[source]#

Load model’s weights from checkpoint at filename.

Differently from load_from_checkpoint(), this method allows to load the state_dict also for models instantiated outside the predictor, without checking that hyperparameters of the checkpoint’s model are the same of the predictor’s model.

predict_batch(batch: Data, preprocess: bool = False, postprocess: bool = True, return_target: bool = False, **forward_kwargs)[source]#

This method takes as input a Data object and outputs the predictions.

Note that this method works seamlessly for all Data subclasses like StaticBatch and DisjointBatch.

Parameters:

batch (Data) – The batch to be forwarded to the model.
preprocess (bool, optional) – If True, then preprocess tensors in batch.input using transformation modules in batch.transform. Note that inputs are preprocessed before creating the batch by default. (default: False)
postprocess (bool, optional) – If True, then postprocess the model output using transformation modules for batch.target in batch.transform. (default: True)
return_target (bool, optional) – If True, then returns also the prediction target batch.target and the prediction mask batch.mask, besides the model output. In this case, the order of the arguments in the return is batch.target, y_hat, batch.mask. (default: False)
**forward_kwargs – additional keyword arguments passed to the forward method.

collate_prediction_outputs(outputs)[source]#

Collate the outputs of the predict_step() method.

Parameters:: outputs – Collated outputs of the predict_step() method.
Returns:: The collated outputs.

class Imputer(model: Optional[Module] = None, loss_fn: Optional[Callable] = None, scale_target: bool = False, metrics: Optional[Mapping[str, Metric]] = None, *, whiten_prob: Optional[Union[float, List[float]]] = 0.05, prediction_loss_weight: float = 1.0, impute_only_missing: bool = True, warm_up_steps: Union[int, Tuple[int, int]] = 0, model_class: Optional[Type] = None, model_kwargs: Optional[Mapping] = None, optim_class: Optional[Type] = None, optim_kwargs: Optional[Mapping] = None, scheduler_class: Optional = None, scheduler_kwargs: Optional[Mapping] = None)[source]#

LightningModule to implement imputers.

An imputer is an engines designed to fill out missing values in spatiotemporal data.

Parameters:

model (torch.nn.Module, optional) – Model implementing the imputer. Ignored if argument model_class is not null. This argument should mainly be used for inference. (default: None)
model_class (type, optional) – Class of Module implementing the imputer. If not None, argument model will be ignored. (default: None)
model_kwargs (mapping, optional) – Dictionary of arguments to be forwarded to model_class at instantiation. (default: None)
optim_class (type, optional) – Class of Optimizer implementing the optimizer to be used for training the model. (default: None)
optim_kwargs (mapping, optional) – Dictionary of arguments to be forwarded to optim_class at instantiation. (default: None)
loss_fn (callable, optional) – Loss function to be used for training the model. (default: None)
scale_target (bool) – Whether to scale target before evaluating the loss. The metrics instead will always be evaluated in the original range. (default: False)
whiten_prob (float or list) – Randomly mask out a valid datapoint during a training step with probability whiten_prob. If a list is passed, whiten_prob is sampled from the list for each batch. (default: 0.05)
prediction_loss_weight (float) –
The weight to assign to predictions (if any) in the loss. The loss is computed as

\[L = \ell(\bar{y}, y, m) + \lambda \sum_i \ell(\hat{y}_i, y, m)\]

where \(\ell(\bar{y}, y, m)\) is the imputation loss, \(\ell(\bar{y}_i, y, m)\) is the forecasting error of prediction \(\bar{y}_i\), and \(\lambda\) is prediction_loss_weight. (default: 1.0)
impute_only_missing (bool) – Whether to impute only missing values in inference or the whole sequence. (default: True)
warm_up_steps (int, tuple) – Number of steps to be considered as warm up stage at the beginning of the sequence. If a tuple is provided, the padding is applied both at the beginning and the end of the sequence. (default: 0)
metrics (mapping, optional) – Set of metrics to be logged during train, val and test steps. The metric’s name will be automatically prefixed with the loop in which the metric is computed (e.g., metric mae will be logged as train_mae when evaluated during training). (default: None)
scheduler_class (type) – Class of _LRScheduler implementing the learning rate scheduler to be used during training. (default: None)
scheduler_kwargs (mapping) – Dictionary of arguments to be forwarded to scheduler_class at instantiation. (default: None)

trim_warm_up(*args)[source]#: Trim all tensors in args removing a number of first and last steps equals to (self.warm_up_steps[0], self.warm_up_steps[1]), respectively.

on_train_batch_start(batch, batch_idx: int) → None[source]#: For every training batch, randomly mask out value with probability p = self.whiten_prob. Then, whiten missing values in batch.input.x.