pipeline package#

This package handles the various steps of the GNN-based pipeline, from preprocessing to track-finding evaluation.

pipeline.get_model(path_or_config, step)[source]#

Get the model class of a given step for a given pipeline configuration.

Parameters:

  • path_or_config (str | dict) – pipeline configuration

  • step (str) – model step (embedding or gnn)

Return type:

Type[ModelBase]

Returns:

Model class that can be later instantiated.

pipeline.instantiate_model_for_training(path_or_config, step)[source]#

Instantiate a new model. The model can then be trained.

The function auto-detects the the model type. The latter can also be instantiated from a trained model (transfer learning) using the parameter from in the configuration file.

Parameters:

  • path_or_config (str | dict) – pipeline configuraiton

  • step (str) – model step (embedding or gnn)

Return type:

ModelBase

Returns:

Model that was not trained.

pipeline.load_trained_model(path_or_config, step, checkpoint_path=None, **kwargs)[source]#

Load a model that was already trained.

Parameters:

  • path_or_config (str | dict) – pipeline configuration

  • step (str) – model step (embedding or gnn)

  • checkpoint_path (Optional[str]) – path to a checkpoint

Return type:

ModelBase

Returns:

Trained model

Subpackages#