sleap.nn.training

Contents

sleap.nn.training#

Training functionality and high level APIs.

class sleap.nn.training.BottomUpModelTrainer(data_readers: DataReaders, model: Model, config: TrainingJobConfig, initial_config: TrainingJobConfig | None = None)[source]#

Trainer for models that output multi-instance confidence maps and PAFs.

property has_offsets: bool#

Whether model is configured to output refinement offsets.

property input_keys: List[str]#

Return example keys to be mapped to model inputs.

property output_keys: List[str]#

Return example keys to be mapped to model outputs.

class sleap.nn.training.BottomUpMultiClassModelTrainer(data_readers: DataReaders, model: Model, config: TrainingJobConfig, initial_config: TrainingJobConfig | None = None)[source]#

Trainer for models that output multi-instance confidence maps and class maps.

property has_offsets: bool#

Whether model is configured to output refinement offsets.

property input_keys: List[str]#

Return example keys to be mapped to model inputs.

property output_keys: List[str]#

Return example keys to be mapped to model outputs.

class sleap.nn.training.CentroidConfmapsModelTrainer(data_readers: DataReaders, model: Model, config: TrainingJobConfig, initial_config: TrainingJobConfig | None = None)[source]#

Trainer for models that output centroid confidence maps.

property has_offsets: bool#

Whether model is configured to output refinement offsets.

property input_keys: List[str]#

Return example keys to be mapped to model inputs.

property output_keys: List[str]#

Return example keys to be mapped to model outputs.

class sleap.nn.training.DataReaders(training_labels_reader: LabelsReader, validation_labels_reader: LabelsReader, test_labels_reader: LabelsReader | None = None)[source]#

Container class for SLEAP labels that serve as training data sources.

training_labels_reader#

LabelsReader pipeline provider for a training data from a Labels instance.

Type:

sleap.nn.data.providers.LabelsReader

validation_labels_reader#

LabelsReader pipeline provider for a validation data from a Labels instance.

Type:

sleap.nn.data.providers.LabelsReader

test_labels_reader#

LabelsReader pipeline provider for a test set data from a Labels instance. This is not necessary for training.

Type:

sleap.nn.data.providers.LabelsReader | None

classmethod from_config(labels_config: LabelsConfig, training: str | Labels, validation: str | Labels | float, test: str | Labels | None = None, video_search_paths: List[str] | None = None, update_config: bool = False, with_track_only: bool = False) DataReaders[source]#

Create data readers from a (possibly incomplete) configuration.

classmethod from_labels(training: str | Labels, validation: str | Labels | float, test: str | Labels | None = None, video_search_paths: List[str] | None = None, labels_config: LabelsConfig | None = None, update_config: bool = False, with_track_only: bool = False) DataReaders[source]#

Create data readers from Labels datasets as data providers.

property test_labels: Labels#

Return the Labels underlying the test data reader.

property training_labels: Labels#

Return the Labels underlying the training data reader.

property validation_labels: Labels#

Return the Labels underlying the validation data reader.

class sleap.nn.training.SingleInstanceModelTrainer(data_readers: DataReaders, model: Model, config: TrainingJobConfig, initial_config: TrainingJobConfig | None = None)[source]#

Trainer for models that output single-instance confidence maps.

property has_offsets: bool#

Whether model is configured to output refinement offsets.

property input_keys: List[str]#

Return example keys to be mapped to model inputs.

property output_keys: List[str]#

Return example keys to be mapped to model outputs.

class sleap.nn.training.TopDownMultiClassModelTrainer(data_readers: DataReaders, model: Model, config: TrainingJobConfig, initial_config: TrainingJobConfig | None = None)[source]#

Trainer for models that output multi-instance confidence maps and class maps.

property has_offsets: bool#

Whether model is configured to output refinement offsets.

property input_keys: List[str]#

Return example keys to be mapped to model inputs.

property output_keys: List[str]#

Return example keys to be mapped to model outputs.

class sleap.nn.training.TopdownConfmapsModelTrainer(data_readers: DataReaders, model: Model, config: TrainingJobConfig, initial_config: TrainingJobConfig | None = None)[source]#

Trainer for models that output instance centered confidence maps.

property has_offsets: bool#

Whether model is configured to output refinement offsets.

property input_keys: List[str]#

Return example keys to be mapped to model inputs.

property output_keys: List[str]#

Return example keys to be mapped to model outputs.

class sleap.nn.training.Trainer(data_readers: DataReaders, model: Model, config: TrainingJobConfig, initial_config: TrainingJobConfig | None = None)[source]#

Base trainer class that provides general model training functionality.

This class is intended to be instantiated using the from_config() class method, which will return the appropriate subclass based on the input configuration.

This class should not be used directly. It is intended to be subclassed by a model output type-specific trainer that provides more specific functionality.

data_readers#

A DataReaders instance that contains training data providers.

Type:

sleap.nn.training.DataReaders

model#

A Model instance describing the SLEAP model to train.

Type:

sleap.nn.model.Model

config#

A TrainingJobConfig that describes the training parameters.

Type:

sleap.nn.config.training_job.TrainingJobConfig

initial_config#

This attribute will contain a copy of the input configuration before any attributes are updated in config.

Type:

sleap.nn.config.training_job.TrainingJobConfig | None

pipeline_builder#

A model output type-specific data pipeline builder to create pipelines that generate data used for training. This must be specified in subclasses.

Type:

sleap.nn.training.PipelineBuilder

training_pipeline#

The data pipeline that generates examples from the training set for optimization.

Type:

sleap.nn.data.pipelines.Pipeline

validation_pipeline#

The data pipeline that generates examples from the validation set for optimization.

Type:

sleap.nn.data.pipelines.Pipeline

training_viz_pipeline#

The data pipeline that generates examples from the training set for visualization.

Type:

sleap.nn.data.pipelines.Pipeline

validation_viz_pipeline#

The data pipeline that generates examples from the validation set for visualization.

Type:

sleap.nn.data.pipelines.Pipeline

optimization_callbacks#

Keras callbacks related to optimization.

Type:

List[keras.callbacks.Callback]

output_callbacks#

Keras callbacks related to outputs.

Type:

List[keras.callbacks.Callback]

visualization_callbacks#

Keras callbacks related to visualization.

Type:

List[keras.callbacks.Callback]

run_path#

The path to the run folder that will contain training results, if any.

Type:

str | None

property callbacks: List[Callback]#

Return all callbacks currently configured.

cleanup()[source]#

Delete visualization images subdirectory.

evaluate()[source]#

Compute evaluation metrics on data splits and save them.

classmethod from_config(config: TrainingJobConfig, training_labels: str | Labels | None = None, validation_labels: str | Labels | float | None = None, test_labels: str | Labels | None = None, video_search_paths: List[str] | None = None) Trainer[source]#

Initialize the trainer from a training job configuration.

Parameters:
  • config – A TrainingJobConfig instance.

  • training_labels – Training labels to use instead of the ones in the config, if any. If a path is specified, it will overwrite the one in the config.

  • validation_labels – Validation labels to use instead of the ones in the config, if any. If a path is specified, it will overwrite the one in the config.

  • test_labels – Teset labels to use instead of the ones in the config, if any. If a path is specified, it will overwrite the one in the config.

abstract property input_keys#

Implement in subclasses.

property keras_model: Model#

Alias for self.model.keras_model.

abstract property output_keys#

Implement in subclasses.

package()[source]#

Package model folder into a zip file for portability.

setup()[source]#

Set up data pipeline and model for training.

train()[source]#

Execute the optimization loop to train the model.

sleap.nn.training.create_trainer_using_cli(args: List | None = None)[source]#

Create CLI for training.

sleap.nn.training.get_timestamp() str[source]#

Return the date and time as a string.

sleap.nn.training.main(args: List | None = None)[source]#

Create CLI for training and run.

sleap.nn.training.sanitize_scope_name(name: str) str[source]#

Sanitizes string which will be used as TensorFlow scope name.

sleap.nn.training.setup_checkpointing(config: CheckpointingConfig, run_path: str) List[Callback][source]#

Set up model checkpointing callbacks from config.

sleap.nn.training.setup_losses(config: OptimizationConfig) Callable[[Tensor, Tensor], Tensor][source]#

Set up model loss function from config.

sleap.nn.training.setup_metrics(config: OptimizationConfig, part_names: List[str] | None = None) List[Loss | Metric][source]#

Set up training metrics from config.

sleap.nn.training.setup_new_run_folder(config: OutputsConfig, base_run_name: str | None = None) str[source]#

Create a new run folder from config.

sleap.nn.training.setup_optimization_callbacks(config: OptimizationConfig) List[Callback][source]#

Set up optimization callbacks from config.

sleap.nn.training.setup_optimizer(config: OptimizationConfig) OptimizerV2[source]#

Set up model optimizer from config.

sleap.nn.training.setup_output_callbacks(config: OutputsConfig, run_path: str | None = None) List[Callback][source]#

Set up training outputs callbacks from config.

sleap.nn.training.setup_tensorboard(config: TensorBoardConfig, run_path: str) List[Callback][source]#

Set up TensorBoard callbacks from config.

sleap.nn.training.setup_visualization(config: OutputsConfig, run_path: str, viz_fn: Callable[[], Figure], name: str) List[Callback][source]#

Set up visualization callbacks from config.

sleap.nn.training.setup_zmq_callbacks(zmq_config: ZMQConfig) List[Callback][source]#

Set up ZeroMQ callbacks from config.