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.

Checking a Config First¶

Add --dry-run to check a config before a long run:

sleap-nn train --config config.yaml --dry-run

It loads the labels, resolves and validates the config, builds the datasets and model, and runs one training batch and one validation batch on a single device. It uses the config's own settings, including its data pipeline and dataloader workers, but builds the datasets from just one batch's worth of frames, so even a cached pipeline takes seconds. Then it prints a summary (model type, frame and sample counts and steps per epoch for the full data, parameter count, run folder, data pipeline, whether the image cache fits in memory) and exits. Config and data errors show up in seconds instead of after training has started.

It also checks that the run can write where it will: it writes the configs and a checkpoint (and any disk image cache) to a hidden .dryrun-* folder in ckpt_dir (and in cache_img_path, if set), then removes them, along with any of those folders it had to create. A dry run that is killed partway can leave a .dryrun-* folder behind, which is safe to delete.

A dry run uses one device, so it doesn't check that multi-GPU setup works.

To resume an interrupted run, see Resume & Fine-Tune.

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 (and, optionally, validation data)
sleap-nn train --config config.yaml --labels train.slp
sleap-nn train --config config.yaml --labels train.slp --val-labels val.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 --labels labels.pkg.slp

# Train centered instance model
sleap-nn train -d /path/to/configs -c centered_instance_unet --labels 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