sleap.nn.training

Training functionality and high level APIs.

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

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

property input_keys

Return example keys to be mapped to model inputs.

property output_keys

Return example keys to be mapped to model outputs.

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

Trainer for models that output centroid confidence maps.

property input_keys

Return example keys to be mapped to model inputs.

property output_keys

Return example keys to be mapped to model outputs.

class sleap.nn.training.DataReaders(training_labels_reader: sleap.nn.data.providers.LabelsReader, validation_labels_reader: sleap.nn.data.providers.LabelsReader, test_labels_reader: Optional[sleap.nn.data.providers.LabelsReader] = None)

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

training_labels_reader

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

validation_labels_reader

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

test_labels_reader

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

classmethod from_config(labels_config: sleap.nn.config.data.LabelsConfig, training: Union[str, sleap.io.dataset.Labels], validation: Union[str, sleap.io.dataset.Labels, float], test: Optional[Union[str, sleap.io.dataset.Labels]] = None, video_search_paths: Optional[List[str]] = None, update_config: bool = False) → sleap.nn.training.DataReaders

Create data readers from a (possibly incomplete) configuration.

classmethod from_labels(training: Union[str, sleap.io.dataset.Labels], validation: Union[str, sleap.io.dataset.Labels, float], test: Optional[Union[str, sleap.io.dataset.Labels]] = None, video_search_paths: Optional[List[str]] = None) → sleap.nn.training.DataReaders

Create data readers from sleap.Labels datasets as data providers.

property test_labels

Return the sleap.Labels underlying the test data reader.

property training_labels

Return the sleap.Labels underlying the training data reader.

property validation_labels

Return the sleap.Labels underlying the validation data reader.

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

Trainer for models that output single-instance confidence maps.

property input_keys

Return example keys to be mapped to model inputs.

property output_keys

Return example keys to be mapped to model outputs.

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

Trainer for models that output instance centered confidence maps.

property input_keys

Return example keys to be mapped to model inputs.

property output_keys

Return example keys to be mapped to model outputs.

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

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.

model

A Model instance describing the SLEAP model to train.

config

A TrainingJobConfig that describes the training parameters.

initial_config

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

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.

training_pipeline

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

validation_pipeline

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

training_viz_pipeline

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

validation_viz_pipeline

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

optimization_callbacks

Keras callbacks related to optimization.

output_callbacks

Keras callbacks related to outputs.

visualization_callbacks

Keras callbacks related to visualization.

run_path

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

property callbacks

Return all callbacks currently configured.

classmethod from_config(config: sleap.nn.config.training_job.TrainingJobConfig, training_labels: Optional[Union[str, sleap.io.dataset.Labels]] = None, validation_labels: Optional[Union[str, sleap.io.dataset.Labels, float]] = None, test_labels: Optional[Union[str, sleap.io.dataset.Labels]] = None, video_search_paths: Optional[List[str]] = None) → sleap.nn.training.Trainer

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

Alias for self.model.keras_model.

abstract property output_keys

Implement in subclasses.

setup()

Set up data pipeline and model for training.

train()

Execute the optimization loop to train the model.

sleap.nn.training.get_timestamp() → str

Return the date and time as a string.

sleap.nn.training.main()

Create CLI for training and run.

sleap.nn.training.sanitize_scope_name(name: str) → str

Sanitizes string which will be used as TensorFlow scope name.

sleap.nn.training.setup_checkpointing(config: sleap.nn.config.outputs.CheckpointingConfig, run_path: str) → List[tensorflow.python.keras.callbacks.Callback]

Set up model checkpointing callbacks from config.

sleap.nn.training.setup_losses(config: sleap.nn.config.optimization.OptimizationConfig) → Callable[tensorflow.python.framework.ops.Tensor, tensorflow.python.framework.ops.Tensor]

Set up model loss function from config.

sleap.nn.training.setup_metrics(config: sleap.nn.config.optimization.OptimizationConfig, part_names: Optional[List[str]] = None) → List[Union[tensorflow.python.keras.losses.Loss, tensorflow.python.keras.metrics.Metric]]

Set up training metrics from config.

sleap.nn.training.setup_new_run_folder(config: sleap.nn.config.outputs.OutputsConfig, base_run_name: Optional[str] = None) → str

Create a new run folder from config.

sleap.nn.training.setup_optimization_callbacks(config: sleap.nn.config.optimization.OptimizationConfig) → List[tensorflow.python.keras.callbacks.Callback]

Set up optimization callbacks from config.

sleap.nn.training.setup_optimizer(config: sleap.nn.config.optimization.OptimizationConfig) → tensorflow.python.keras.optimizer_v2.optimizer_v2.OptimizerV2

Set up model optimizer from config.

sleap.nn.training.setup_output_callbacks(config: sleap.nn.config.outputs.OutputsConfig, run_path: Optional[str] = None) → List[tensorflow.python.keras.callbacks.Callback]

Set up training outputs callbacks from config.

sleap.nn.training.setup_tensorboard(config: sleap.nn.config.outputs.TensorBoardConfig, run_path: str) → List[tensorflow.python.keras.callbacks.Callback]

Set up TensorBoard callbacks from config.

sleap.nn.training.setup_visualization(config: sleap.nn.config.outputs.OutputsConfig, run_path: str, viz_fn: Callable[matplotlib.figure.Figure], name: str) → List[tensorflow.python.keras.callbacks.Callback]

Set up visualization callbacks from config.

sleap.nn.training.setup_zmq_callbacks(zmq_config: sleap.nn.config.outputs.ZMQConfig) → List[tensorflow.python.keras.callbacks.Callback]

Set up ZeroMQ callbacks from config.