Training Models¶

Train pose estimation models with SLEAP-NN.

New to SLEAP-NN?

See Model Types to understand the different model architectures (single instance, top-down, bottom-up) and when to use each.

Using uv workflow

If using uvx, no installation needed
If using uv sync, prefix commands with uv run:
```
uv run sleap-nn train ...
```

Basic Training¶

Using CLI¶

sleap-nn train --config config.yaml

Or with separate config directory and name:

sleap-nn train --config-dir /path/to/configs --config-name my_config

See the CLI Reference for all available parameters.

Using Python API¶

from omegaconf import OmegaConf
from sleap_nn.train import run_training

config = OmegaConf.load("config.yaml")
run_training(config=config)

With Custom Labels¶

import sleap_io as sio
from sleap_nn.train import run_training

config = OmegaConf.load("config.yaml")
train_labels = sio.load_slp("train.slp")
val_labels = sio.load_slp("val.slp")

run_training(config=config,
            train_labels=[train_labels],
            val_labels=[val_labels])

Config Overrides¶

Override any config value from the command line:

# Change epochs
sleap-nn train --config config.yaml trainer_config.max_epochs=200

# Change learning rate
sleap-nn train --config config.yaml trainer_config.optimizer.lr=0.0005

# Change batch size
sleap-nn train --config config.yaml trainer_config.train_data_loader.batch_size=8

# Set training data
sleap-nn train --config config.yaml "data_config.train_labels_path=[train.slp]"

# Set number of GPUs
sleap-nn train --config config.yaml trainer_config.trainer_devices=1

Video Path Remapping¶

When training on a different machine than where labels were created:

Replace by orderMap specific pathsReplace prefix

sleap-nn train --config config.yaml \
    --video-paths /new/path/video1.mp4 \
    --video-paths /new/path/video2.mp4

sleap-nn train --config config.yaml \
    --video-path-map /old/video.mp4 /new/video.mp4

sleap-nn train --config config.yaml \
    --prefix-map /old/server/data /new/local/data

Choose one option

You can only use one of --video-paths, --video-path-map, or --prefix-map at a time.

Training Without Config¶

Quick training with minimal setup using presets:

from sleap_nn.train import train

train(
    train_labels_path=["labels.slp"],
    backbone_config="unet_medium_rf",  # or "unet_large_rf"
    head_configs="bottomup",           # or "single_instance", etc.
    save_ckpt=True,
)

With augmentation:

train(
    train_labels_path=["labels.slp"],
    backbone_config="unet_medium_rf",
    head_configs="bottomup",
    use_augmentations_train=True,
    intensity_aug="uniform_noise",
    geometric_aug=["rotation", "scale"],
)

Top-Down Training¶

Top-down models need two separate training runs. See Model Types for details on when to use top-down vs bottom-up.

# Train centroid model
sleap-nn train -d /path/to/configs -c centroid_unet \
    "data_config.train_labels_path=[labels.pkg.slp]"

# Train centered instance model
sleap-nn train -d /path/to/configs -c centered_instance_unet \
    "data_config.train_labels_path=[labels.pkg.slp]"

Monitoring Training¶

Track training progress with WandB logging, visualizations, and evaluation metrics.

trainer_config:
  use_wandb: true
  wandb:
    entity: your-username
    project: your-project
    viz_enabled: true   # Log prediction visualizations

  eval:
    enabled: true       # Compute evaluation metrics during training
    frequency: 1        # Evaluate every epoch

For detailed configuration options including:

WandB visualization settings (interactive boxes, confidence map masks)
Epoch-end evaluation metrics (OKS, PCK, centroid metrics)
Local visualization output
Per-head loss monitoring

See the dedicated guide:

Monitoring & Visualization Guide

Checkpointing & Artifacts¶

Each training run creates a checkpoint directory with:

File	Description
`best.ckpt`	Best model weights
`initial_config.yaml`	Original user config
`training_config.yaml`	Full config with computed values
`labels_gt.train.0.slp`	Training data split (ground truth)
`labels_gt.val.0.slp`	Validation data split (ground truth)
`labels_pr.train.slp`	Predictions on training data
`labels_pr.val.slp`	Predictions on validation data
`metrics.train.0.npz`	Training metrics
`metrics.val.0.npz`	Validation metrics
`training_log.csv`	Loss/metrics per epoch

Resuming & Fine-Tuning¶

To resume an interrupted run or fine-tune from pre-trained weights, see the dedicated guide:

Resuming & Fine-Tuning Guide

Multi-GPU Training¶

For multi-GPU and distributed training, see the dedicated guide:

Multi-GPU Training Guide

Performance Tips¶

Enable Caching¶

data_config:
  data_pipeline_fw: torch_dataset_cache_img_memory  # RAM caching
  # or
  data_pipeline_fw: torch_dataset_cache_img_disk    # Disk caching

With caching, you can use num_workers > 0:

trainer_config:
  train_data_loader:
    num_workers: 4

Workers without caching

Keep num_workers: 0 when not using caching.

Understanding Training¶

Epochs and Batches¶

Training occurs in epochs, where one epoch consists of the larger of:

(number of training images) / (batch size), or
200 batches

With larger datasets, one epoch equals one pass over the training data.

Early Stopping¶

By default, training stops early when a plateau is detected in the validation loss to prevent overfitting. You can disable this or set a fixed number of epochs:

trainer_config:
  max_epochs: 200
  early_stopping:
    stop_training_on_plateau: false  # Disable early stopping

Augmentation Strategy¶

During training, augmentations are applied to raw images and poses to generate variants of labeled data. This promotes generalization.

Rotation recommendations:

Overhead/top-down view: Use full rotation range (-180° to 180°)
Side view: Use limited rotation (-15° to 15°)

data_config:
  augmentation_config:
    geometric:
      rotation_min: -180.0
      rotation_max: 180.0