Training and inference on an example dataset¶

In this notebook we'll install sleap-nn, download a sample dataset, run training and inference on that dataset using the sleap-nn, and then download the predictions.

Note: If you only need to perform training and inference (and not use the full SLEAP GUI or labeling tools), you don't need to install the entire sleap package. You can just install sleap-nn, which is a lighter-weight package focused on model training and inference.

Install sleap-nn¶

In [ ]:

Copied!

!pip install -qqq "sleap-nn[torch-cpu]"

# if you have GPU (in colab, enable GPU runtime)
# !pip install -qqq "sleap-nn[torch-cuda128]"
!pip install -qqq "sleap-nn[torch-cpu]"

# if you have GPU (in colab, enable GPU runtime)
# !pip install -qqq "sleap-nn[torch-cuda128]"

zsh:1: command not found: pip

Download sample training data into Colab¶

Let's download a sample dataset from the SLEAP sample datasets repository into Colab.

In [2]:

Copied!





!apt-get install tree
!wget -O dataset.zip https://github.com/talmolab/sleap-datasets/releases/download/dm-courtship-v1/drosophila-melanogaster-courtship.zip
!mkdir dataset
!unzip dataset.zip -d dataset
!rm dataset.zip
!tree dataset
!apt-get install tree
!wget -O dataset.zip https://github.com/talmolab/sleap-datasets/releases/download/dm-courtship-v1/drosophila-melanogaster-courtship.zip
!mkdir dataset
!unzip dataset.zip -d dataset
!rm dataset.zip
!tree dataset

zsh:1: command not found: apt-get
--2025-09-23 22:49:59--  https://github.com/talmolab/sleap-datasets/releases/download/dm-courtship-v1/drosophila-melanogaster-courtship.zip
Resolving github.com (github.com)... 140.82.116.4
Connecting to github.com (github.com)|140.82.116.4|:443... connected.
HTTP request sent, awaiting response... 302 Found
Location: https://release-assets.githubusercontent.com/github-production-release-asset/263375180/16df8d00-94f1-11ea-98d1-6c03a2f89e1c?sp=r&sv=2018-11-09&sr=b&spr=https&se=2025-09-24T06%3A32%3A23Z&rscd=attachment%3B+filename%3Ddrosophila-melanogaster-courtship.zip&rsct=application%2Foctet-stream&skoid=96c2d410-5711-43a1-aedd-ab1947aa7ab0&sktid=398a6654-997b-47e9-b12b-9515b896b4de&skt=2025-09-24T05%3A32%3A15Z&ske=2025-09-24T06%3A32%3A23Z&sks=b&skv=2018-11-09&sig=Ij3ERUXEA5fEIAbcmukQghtno0Fl4j0%2BrI9epJ%2FH4Jw%3D&jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmVsZWFzZS1hc3NldHMuZ2l0aHVidXNlcmNvbnRlbnQuY29tIiwia2V5Ijoia2V5MSIsImV4cCI6MTc1ODY5MzI5OSwibmJmIjoxNzU4NjkyOTk5LCJwYXRoIjoicmVsZWFzZWFzc2V0cHJvZHVjdGlvbi5ibG9iLmNvcmUud2luZG93cy5uZXQifQ.Z7KygNVvJl9-7hx2zZxvJ2Nk11iVAFbZHyRtvwuRNOA&response-content-disposition=attachment%3B%20filename%3Ddrosophila-melanogaster-courtship.zip&response-content-type=application%2Foctet-stream [following]
--2025-09-23 22:49:59--  https://release-assets.githubusercontent.com/github-production-release-asset/263375180/16df8d00-94f1-11ea-98d1-6c03a2f89e1c?sp=r&sv=2018-11-09&sr=b&spr=https&se=2025-09-24T06%3A32%3A23Z&rscd=attachment%3B+filename%3Ddrosophila-melanogaster-courtship.zip&rsct=application%2Foctet-stream&skoid=96c2d410-5711-43a1-aedd-ab1947aa7ab0&sktid=398a6654-997b-47e9-b12b-9515b896b4de&skt=2025-09-24T05%3A32%3A15Z&ske=2025-09-24T06%3A32%3A23Z&sks=b&skv=2018-11-09&sig=Ij3ERUXEA5fEIAbcmukQghtno0Fl4j0%2BrI9epJ%2FH4Jw%3D&jwt=eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmVsZWFzZS1hc3NldHMuZ2l0aHVidXNlcmNvbnRlbnQuY29tIiwia2V5Ijoia2V5MSIsImV4cCI6MTc1ODY5MzI5OSwibmJmIjoxNzU4NjkyOTk5LCJwYXRoIjoicmVsZWFzZWFzc2V0cHJvZHVjdGlvbi5ibG9iLmNvcmUud2luZG93cy5uZXQifQ.Z7KygNVvJl9-7hx2zZxvJ2Nk11iVAFbZHyRtvwuRNOA&response-content-disposition=attachment%3B%20filename%3Ddrosophila-melanogaster-courtship.zip&response-content-type=application%2Foctet-stream
Resolving release-assets.githubusercontent.com (release-assets.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.111.133, ...
Connecting to release-assets.githubusercontent.com (release-assets.githubusercontent.com)|185.199.109.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 111973079 (107M) [application/octet-stream]
Saving to: ‘dataset.zip’

dataset.zip         100%[===================>] 106.79M  39.0MB/s    in 2.7s    

2025-09-23 22:50:02 (39.0 MB/s) - ‘dataset.zip’ saved [111973079/111973079]

Archive:  dataset.zip
   creating: dataset/drosophila-melanogaster-courtship
  inflating: dataset/drosophila-melanogaster-courtship/.DS_Store  
   creating: dataset/__MACOSX
   creating: dataset/__MACOSX/drosophila-melanogaster-courtship
  inflating: dataset/__MACOSX/drosophila-melanogaster-courtship/._.DS_Store  
  inflating: dataset/drosophila-melanogaster-courtship/20190128_113421.mp4  
  inflating: dataset/__MACOSX/drosophila-melanogaster-courtship/._20190128_113421.mp4  
  inflating: dataset/drosophila-melanogaster-courtship/courtship_labels.slp  
  inflating: dataset/__MACOSX/drosophila-melanogaster-courtship/._courtship_labels.slp  
  inflating: dataset/drosophila-melanogaster-courtship/example.jpg  
  inflating: dataset/__MACOSX/drosophila-melanogaster-courtship/._example.jpg  
zsh:1: command not found: tree

Train models¶

For the top-down pipeline, we'll need train two models: a centroid model and a centered-instance model.

We'll first train a model for centroids using the default training profile. The training profile determines the model architecture, the learning rate, and other parameters.

When you start training, you'll first see the training parameters and then the training and validation loss for each training epoch.

As soon as you're satisfied with the validation loss you see for an epoch during training, you're welcome to stop training by clicking the stop button. The version of the model with the lowest validation loss is saved during training, and that's what will be used for inference.

If you don't stop training, it will run for 200 epochs or until validation loss fails to improve for some number of epochs (controlled by the early_stopping fields in the training profile).

In [3]:

Copied!

# Let's get the default config files
!wget -O baseline.centroid.yaml https://raw.githubusercontent.com/talmolab/sleap-nn/refs/heads/main/docs/sample_configs/config_centroid_unet.yaml
!wget -O baseline.centered_instance.yaml https://raw.githubusercontent.com/talmolab/sleap-nn/refs/heads/main/docs/sample_configs/config_topdown_centered_instance_unet_medium_rf.yaml
# Let's get the default config files
!wget -O baseline.centroid.yaml https://raw.githubusercontent.com/talmolab/sleap-nn/refs/heads/main/docs/sample_configs/config_centroid_unet.yaml
!wget -O baseline.centered_instance.yaml https://raw.githubusercontent.com/talmolab/sleap-nn/refs/heads/main/docs/sample_configs/config_topdown_centered_instance_unet_medium_rf.yaml

--2025-09-23 22:50:13--  https://raw.githubusercontent.com/talmolab/sleap-nn/refs/heads/main/docs/sample_configs/config_centroid_unet.yaml
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8000::154, 2606:50c0:8001::154, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 3286 (3.2K) [text/plain]
Saving to: ‘baseline.centroid.yaml’

baseline.centroid.y 100%[===================>]   3.21K  --.-KB/s    in 0s      

2025-09-23 22:50:13 (33.0 MB/s) - ‘baseline.centroid.yaml’ saved [3286/3286]

--2025-09-23 22:50:14--  https://raw.githubusercontent.com/talmolab/sleap-nn/refs/heads/main/docs/sample_configs/config_topdown_centered_instance_unet.yaml
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8000::154, 2606:50c0:8001::154, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 3114 (3.0K) [text/plain]
Saving to: ‘baseline.centered_instance.yaml’

baseline.centered_i 100%[===================>]   3.04K  --.-KB/s    in 0s      

2025-09-23 22:50:14 (32.3 MB/s) - ‘baseline.centered_instance.yaml’ saved [3114/3114]

In [ ]:

Copied!

!sleap-nn train --config baseline.centroid.yaml "data_config.train_labels_path=[dataset/drosophila-melanogaster-courtship/courtship_labels.slp]" trainer_config.ckpt_dir="models" trainer_config.run_name="courtship.centroid" trainer_config.max_epochs=200
!sleap-nn train --config baseline.centroid.yaml "data_config.train_labels_path=[dataset/drosophila-melanogaster-courtship/courtship_labels.slp]" trainer_config.ckpt_dir="models" trainer_config.run_name="courtship.centroid" trainer_config.max_epochs=200

Let's now train a centered-instance model.

In [ ]:

Copied!

!sleap-nn train --config baseline.centered_instance.yaml "data_config.train_labels_path=[dataset/drosophila-melanogaster-courtship/courtship_labels.slp]" trainer_config.ckpt_dir="models" trainer_config.run_name="courtship.topdown_confmaps" trainer_config.max_epochs=200
!sleap-nn train --config baseline.centered_instance.yaml "data_config.train_labels_path=[dataset/drosophila-melanogaster-courtship/courtship_labels.slp]" trainer_config.ckpt_dir="models" trainer_config.run_name="courtship.topdown_confmaps" trainer_config.max_epochs=200

The models (along with the profiles and ground truth data used to train and validate the model) are saved in the models/ directory:

In [11]:

Copied!

!tree models/
!tree models/

models/
├── courtship.centroid
│   ├── best.ckpt
│   ├── initial_config.yaml
│   ├── labels_train_gt_0.slp
│   ├── labels_val_gt_0.slp
│   ├── training_config.yaml
│   └── training_log.csv
└── courtship.topdown_confmaps
    ├── best.ckpt
    ├── initial_config.yaml
    ├── labels_train_gt_0.slp
    ├── labels_val_gt_0.slp
    ├── training_config.yaml
    └── training_log.csv

3 directories, 12 files

Inference¶

Let's run inference with our trained models for centroids and centered instances.

In [9]:

Copied!

!sleap-nn track -i "dataset/drosophila-melanogaster-courtship/20190128_113421.mp4" --frames 0-100 -m "models/courtship.centroid" -m "models/courtship.topdown_confmaps"
!sleap-nn track -i "dataset/drosophila-melanogaster-courtship/20190128_113421.mp4" --frames 0-100 -m "models/courtship.centroid" -m "models/courtship.topdown_confmaps"

2025-09-24 00:12:47 | INFO | sleap_nn.predict:run_inference:319 | Started inference at: 2025-09-24 00:12:47.871203
2025-09-24 00:12:47 | INFO | sleap_nn.predict:run_inference:330 | Integral refinement is not supported with MPS device. Using CPU.
2025-09-24 00:12:47 | INFO | sleap_nn.predict:run_inference:335 | Using device: cpu
Predicting... ━━━━━━━━━━━━━━━ 100% 101/101 ETA: 0:00:00 Elapsed: 0:00:17 6.0 FPS0 FPS7 FPS
2025-09-24 00:13:05 | INFO | sleap_nn.predict:run_inference:453 | Finished inference at: 2025-09-24 00:13:05.531684
2025-09-24 00:13:05 | INFO | sleap_nn.predict:run_inference:454 | Total runtime: 17.660489082336426 secs
2025-09-24 00:13:05 | INFO | sleap_nn.predict:run_inference:465 | Predictions output path: dataset/drosophila-melanogaster-courtship/20190128_113421.predictions.slp
2025-09-24 00:13:05 | INFO | sleap_nn.predict:run_inference:466 | Saved file at: 2025-09-24 00:13:05.594857

When inference is finished, predictions are saved in a file. Since we didn't specify a path, it will be saved as <video filename>.predictions.slp in the same directory as the video:

In [10]:

Copied!

!tree dataset/drosophila-melanogaster-courtship
!tree dataset/drosophila-melanogaster-courtship

dataset/drosophila-melanogaster-courtship
├── 20190128_113421.mp4
├── 20190128_113421.mp4.predictions.slp
├── 20190128_113421.predictions.slp
├── courtship_labels.slp
└── example.jpg

1 directory, 5 files

If you're using Chrome you can download your trained models like so:

In [31]:

Copied!





# Zip up the models directory
!zip -r trained_models.zip models/

# Download.
from google.colab import files
files.download("/content/trained_models.zip")
# Zip up the models directory
!zip -r trained_models.zip models/

# Download.
from google.colab import files
files.download("/content/trained_models.zip")

  adding: models/ (stored 0%)
  adding: models/courtship.topdown_confmaps/ (stored 0%)
  adding: models/courtship.topdown_confmaps/labels_pr.val.slp (deflated 74%)
  adding: models/courtship.topdown_confmaps/metrics.val.npz (deflated 0%)
  adding: models/courtship.topdown_confmaps/labels_pr.train.slp (deflated 67%)
  adding: models/courtship.topdown_confmaps/labels_gt.val.slp (deflated 72%)
  adding: models/courtship.topdown_confmaps/initial_config.json (deflated 73%)
  adding: models/courtship.topdown_confmaps/training_log.csv (deflated 55%)
  adding: models/courtship.topdown_confmaps/metrics.train.npz (deflated 0%)
  adding: models/courtship.topdown_confmaps/labels_gt.train.slp (deflated 61%)
  adding: models/courtship.topdown_confmaps/best_model.h5 (deflated 8%)
  adding: models/courtship.topdown_confmaps/training_config.json (deflated 88%)
  adding: models/courtship.centroid/ (stored 0%)
  adding: models/courtship.centroid/labels_pr.val.slp (deflated 82%)
  adding: models/courtship.centroid/metrics.val.npz (deflated 1%)
  adding: models/courtship.centroid/labels_pr.train.slp (deflated 79%)
  adding: models/courtship.centroid/labels_gt.val.slp (deflated 73%)
  adding: models/courtship.centroid/initial_config.json (deflated 74%)
  adding: models/courtship.centroid/training_log.csv (deflated 57%)
  adding: models/courtship.centroid/metrics.train.npz (deflated 0%)
  adding: models/courtship.centroid/labels_gt.train.slp (deflated 61%)
  adding: models/courtship.centroid/best_model.h5 (deflated 7%)
  adding: models/courtship.centroid/training_config.json (deflated 88%)

And you can likewise download your predictions:

In [33]:

Copied!

from google.colab import files
files.download('dataset/drosophila-melanogaster-courtship/20190128_113421.mp4.predictions.slp')
from google.colab import files
files.download('dataset/drosophila-melanogaster-courtship/20190128_113421.mp4.predictions.slp')

In some other browsers (Safari) you might get an error and you can instead download using the "Files" tab in the side panel (it has a folder icon). Select "Show table of contents" in the "View" menu if you don't see the side panel.