pipeline module

class pipeline.Pipeline(train_dataset, valid_dataset=None, test_dataset=None, model=None, training_params=None, optimizer=None, model_class=None, evaluators=[])[source]

Bases: object

Pipeline class to train and evaluate models. To optimize a model, provide an optimizer and model class. To train a model with fixed parameters, provide a model and training parameters.

Parameters:

  • train_dataset (BaseDataset) – The training dataset.

  • valid_dataset (BaseDataset, optional) – The validation dataset. Default is None.

  • test_dataset (BaseDataset, optional) – The test dataset. Default is None.

  • model (Model, optional) – The model to train. Default is None. If optimizer and model_class are provided, this is not used.

  • training_params (Dict, optional) – The parameters for training the model. Default is None. If optimizer and model_class are provided, this is not used.

  • optimizer (OptunaOptimizer, optional) – The optimizer to use for optimization. Default is None.

  • model_class (Model, optional) – The model class to use for optimization. Default is None.

  • evaluators (List, optional) – The evaluators to use for evaluating the model. Default is [].

Raises:

AssertionError – If neither model and training_params nor optimizer and model_class are provided.

evaluate(dataset)[source]

Evaluate the model on a dataset.

Parameters:

dataset (BaseDataset) – The dataset to evaluate the model on.

Returns:

The metrics evaluated on the dataset.

Return type:

metrics (Dict[str, float])

property model: Model

Get the trained model.

run()[source]

Run the pipeline.

Returns:

The output of the model’s fit method.

Return type:

model_output