Train unsupervised models and evaluate them with your data

What we’ll cover:

• How to train a VaDE-based model for single mouse behaviors

• How to train a Contrastive model for multi-mouse behaviors

• How to visualize and evaluate the quality of your embedding space

• How to plot model performance parameters

This tutorial can be seen as an expansion of the unsupervised turorial and will focus more on the different types of models you can train and which options you have to adjust your taining parameters and preprocessing. As for that other tutorial

You will need GPU pytorch

Hence, make sure it is installed first. How taht is done, is shown right at the beginning of the unsupervised turorial

# If using Google colab, uncomment and run this cell and the one below to set up the environment
# Note: because of how colab handles the installation of local packages, this cell will kill your runtime.
# This is not an error! Just continue with the cells below.
# import os
# !git clone -q https://github.com/mlfpm/deepof.git
# !pip install -q -e deepof --progress-bar off
# os.chdir("deepof")
# !curl --output tutorial_files.zip https://datashare.mpcdf.mpg.de/s/TbSrod2fm5jbHfL/download
# !unzip tutorial_files.zip

# import os
# os.chdir("deepof")
# import os, warnings
# warnings.filterwarnings('ignore')

Now first check if torch finds your GPU

import torch

torch.cuda.is_available()

True

Make sure that this returns True before you continue. NOTE: Even if torch does not find a GPU, the following sections of the tutorial will likely still work, but training actual models will be gruesomly slow.

How to train a VaDE-based model for single mouse behaviors

First, some more imports

import os
import deepof
from matplotlib import pyplot as plt

We load our tutorial project

my_deepof_project = deepof.data.load_project("./tutorial_files/tutorial_project")

Now we will preprocess our data into a dataset that the DeepOF models can work with. For this we apply multiple steps of standardization, scaling and alignment which wil be done in the get_graph_dataset for you. A noteworthy input option are dist_standardize, speed_standardize and coord_standardize.

These denote if the respective datatypes should be standardized as a group (groupwise), which preserves skaling relations between group members but might make it more difficult to achive a stable training result for Vade and VQVAE or, alternatively standardized individually for each column (per_column). A per_column makes training easier, but information about body-relations is lost to a degree, as the per_column-standardization distorts the positions of bodyparts relative to each other.

# This code will generate a dataset using graph representations, as well a some auxiliary objects
preprocessed_object, meta_info, adj_matrix, to_preprocess, global_scaler = my_deepof_project.get_graph_dataset(
    animal_id="B", # Here we train only with one animal, the black mouse
    center="Center",
    align="Spine_1",
    window_size=24,
    window_step=1,
    test_videos=1,
    preprocess=True,
    scale="standard",
    dist_standardize = "groupwise", # can also be "per_column" and None
    speed_standardize = "groupwise",
    coord_standardize = "groupwise",
)
samples_max = None
bin_size = None
bin_index = None
precomputed_bins = None

Loading tables                : 100%|██████████| 5/5 [00:21<00:00,  4.29s/step, step=Get graph info]
Sampling                      : 100%|██████████| 53/53 [00:02<00:00, 19.94table/s]
Scaling                       : 100%|██████████| 53/53 [00:08<00:00,  6.39table/s]
Get training windows          : 100%|██████████| 52/52 [00:00<00:00, 758.91table/s]
Get testing windows           : 100%|██████████| 1/1 [00:00<00:00, 999.12table/s]
Reshaping                     : 100%|██████████| 2/2 [00:01<00:00,  1.51table/s]

The next step would now be to train our model. In this example, we have set the pretrained path to an existing model, which will result in that model simply getting loaded i.e. no new training will occur. If you instead want to train a new model from scratch, you can simply set pretrained to None.

Here we call the train_deepof_model function directly which also get’s run by embedding_model_fitting that can be called on the deepOF project. The main difference is that embedding_model_fitting allows additional binning of the data in advance before model training

import deepof.clustering.training
model_val, model_score, model_part, log_summary = deepof.clustering.training.train_deepof_model(

    # We load a pretrained model. Setting this path will lead to a model getting loaded and all other inputs getting ignored.
    pretrained="./tutorial_files/tutorial_project/Trained_models/models/vade/run_0/best_model_val.pth",

    # Basics
    output_path = os.path.join(my_deepof_project._project_path, my_deepof_project._project_name, "Trained_models"), #where to save the model after training
    preprocessed_object=preprocessed_object, # inputs from your preprocessing
    adjacency_matrix=adj_matrix,
    meta_info=meta_info,
    model_name="Vade",              # Determines the model-type, can be VaDE, VQVAE or Contrastive
    encoder_type="recurrent",       # Determines the type of "funnel" to use to compress your data into the latent space. Can be recurrent, TCN or transformer
    batch_size=1024,
    latent_dim=20,                  # Determins the size of the latent/embedding space
    epochs=50,
    n_clusters=10,
    run=0,                          # run id, change this to change the name of the folder your model will be saved in
    learning_rate=0.0004,           # learning rates for the different parts of the VaDE model
    gmm_learning_rate=0.0001,
    learning_rate_pretrain=0.0018,

    # Losses
    main_clustering_loss=1.0, # Negative, Rewards agreement between posterior encoder distribution and GMM distribution
                              #(i.e. GMM should reflect as good as possible what the encoder is capturing)
    kmeans_loss=0.5, # Pulls embeddings towards cluster centroids i.e. reduces scatter

    reg_cat_clusters=0.0, # Punishes uneven usage of clusters. Should, if activated, only be low

    reg_scatter_weight = 0.0, # Punishes highly noisy clusters (that scatter widely with points all over the place)
    reg_scatter_beta = 0.0,

    repel_weight = 0.0, # Pushes cluster centers away from each other
    repel_length_scale = 0.0,

    nonempty_weight = 0.0167301507709498, # Punishes empty or nearly empty clusters

    temporal_cohesion_weight = 1.62686218712298, # Pushes adjacent frames to be simlar q-wise
    prior_loss_weight=1.0,

    # The same as for pretraining
    nonempty_weight_pretrain = 0.0,
    repel_weight_pretrain = 0.01,
    repel_length_scale_pretrain = 0.0,
    kmeans_loss_pretrain=0.166630457111161,

    # KL modifiers (pushes the model gently towards an useful space with simple geometry)
    kl_annealing_mode="tf_sigmoid",
    kl_warmup=1,
    kl_max_weight=0.57,
    kl_end_weight=1.0,
    kl_cooldown=10,

    # The same as for pretraining (here only encoder variance is regulated)
    kl_warmup_pretrain=1,
    kl_max_weight_pretrain=1.0 ,
    kl_end_weight_pretrain=1.0,
    kl_cooldown_pretrain=10,

    # Init special conditions
    freeze_decoder_epochs=0, # Will deactivate the weight updating for decoder or gmm in the main training loop for the given number of epochs
    freeze_gmm_epochs=0,

    # Teacher model setup (optimizes for linear separability of labels, can be used to anchor vade and to keep it from collapsing)
    use_turtle_teacher=True,
    reinit_gmm_on_refresh=False,
    teacher_gamma=3,
    teacher_head_temp=0.35,
    teacher_task_temp=0.35,
    teacher_alpha_sample_entropy=1.0,
    teacher_outer_steps=500,
    teacher_inner_steps=200,
    include_edges_view=False,
    include_angles_view=False,
    include_nodes_view=True,
    teacher_refresh_every=False,
    teacher_freeze_at=3,
    distill_conf_weight = True,
    distill_conf_thresh=0.6,

    # Lamda modifiers (for teacher distillation, determins influence of Teacher on training per epoch)
    lambda_distill=8.0,
    lambda_decay_start=8,
    lambda_end_weight=0.0,
    lambda_cooldown=12,
)

Using device: cuda
BatchDictDataset: building HDF5 at ./tutorial_files\tutorial_project\Trained_models\Datasets...
  Reason: X shape mismatch
HDF5 built. Samples: 778120, chunks: 1024
HDF5 dataset ready. Samples: 778120, x_shape: (24, 11, 3), a_shape: (24, 11, 1), ang_shape: None
BatchDictDataset: building HDF5 at ./tutorial_files\tutorial_project\Trained_models\Datasets...
  Reason: X shape mismatch
HDF5 built. Samples: 14977, chunks: 1024
HDF5 dataset ready. Samples: 14977, x_shape: (24, 11, 3), a_shape: (24, 11, 1), ang_shape: None
TensorBoard logs -> ./tutorial_files\tutorial_project\Trained_models\logs\vade_run_0
Loading pretrained weights from ./tutorial_files/tutorial_project/Trained_models/models/vade/run_0/best_model_val.pth

========================================================================================
Model info: tutorial_files\tutorial_project\Trained_models\models\vade\run_0\best_model_val_info.txt
----------------------------------------------------------------------------------------
stage: best_val
epoch: 46
train_steps: 35720
val_total: 36.408162943522136

[checkpoint_format]
ckpt_contains: bundle
bundle_keys: state_dict, rebuild_spec, log_summary

[common_cfg]
learning_rate: 0.00036594723903481
model_name: vade
encoder_type: recurrent
batch_size: 1024
latent_dim: 20
epochs: 50
n_components: 10
output_path: ./tutorial_files\tutorial_project\Trained_models
data_path: .
log_history: True
pretrained: None
save_weights: True
run: 0
num_workers: 0
prefetch_factor: 0
use_amp: False
interaction_regularization: 0.000348857818925854
kmeans_loss: 0.5
diag_max_batches: 4
seed: 0
limit_train_batches: 1000
limit_val_batches: 1000

[teacher_cfg]
use_turtle_teacher: True
teacher_gamma: 3
teacher_outer_steps: 500
teacher_inner_steps: 200
teacher_normalize_feats: True
teacher_head_temp: 0.35
teacher_task_temp: 0.35
teacher_alpha_sample_entropy: 1.0
lambda_distill: 8.0
lambda_decay_start: 8
lambda_end_weight: 0.0
lambda_cooldown: 12
distill_sharpen_T: 0.5
distill_conf_weight: True
distill_conf_thresh: 0.6
generic_lambda_distill: 2.0
generic_distill_sharpen_T: 0.5
generic_distill_conf_weight: True
generic_distill_conf_thresh: 0.6
generic_distill_warmup_epochs: 1
distill_class_reweight_beta: 1
distill_class_reweight_cap: 3
include_latent_view: True
include_edges_view: False
include_nodes_view: True
include_angles_view: False
pca_nodes_dim: 32
pca_edges_dim: 32
pca_angles_dim: 32
batch_size_nodes: 4096
batch_size_edges: 8192
batch_size_angles: 8192
teacher_refresh_every: None
teacher_freeze_at: 3
reinit_gmm_on_refresh: False
teacher_batch_size: 2048

[vade_cfg]
learning_rate_pretrain: 0.00183373647106535
gmm_learning_rate: 0.000122350019496003
pretrain_epochs: 10
reg_cat_clusters: 0.0
recluster: False
freeze_gmm_epochs: 0
freeze_decoder_epochs: 0
prior_loss_weight: 1.0
reg_scatter_weight: 0.0
temporal_cohesion_weight: 1.62686218712298
reg_scatter_beta: 0.0
repel_weight: 0.0
repel_length_scale: 0.0
tf_cluster_weight: 1.0
nonempty_weight: 0.0167301507709498
nonempty_p: 2.0
nonempty_floor_percent: 0.05
kmeans_loss_pretrain: 0.166630457111161
repel_weight_pretrain: 0.01
repel_length_scale_pretrain: 0.0
nonempty_weight_pretrain: 0.0
nonempty_p_pretrain: 2.0
nonempty_floor_percent_pretrain: 0.05
kl_annealing_mode: tf_sigmoid
kl_max_weight: 0.567943892625782
kl_warmup: 1
kl_end_weight: 1.0
kl_cooldown: 10
kl_annealing_mode_pretrain: tf_sigmoid
kl_max_weight_pretrain: 1.0
kl_warmup_pretrain: 1
kl_end_weight_pretrain: 1.0
kl_cooldown_pretrain: 10
========================================================================================

Alternatively, we can also just call the model loading function directly. Here we load the same model, but at a different checkpoint. Whilst the validation checkpoint from before is the one at which validation loss was the lowest during training, the score checkpoint denotes the epoch in training with the highest (or almost highest) alignment-score during which clusters were separated with the best combination between high confidence and good balance (i.e. many, well separted and well populated clusters)

import deepof.clustering.model_utils_new

best_path_score="./tutorial_files/tutorial_project/Trained_models/models/vade/run_0/best_model_score.pth"
model_score, log_summary, spec, load_report = deepof.clustering.model_utils_new.load_model_from_ckpt(best_path_score)

========================================================================================
Model info: tutorial_files\tutorial_project\Trained_models\models\vade\run_0\best_model_score_info.txt
----------------------------------------------------------------------------------------
stage: best_score
epoch: 49
train_steps: 38000
val_total: 36.505517578125
score_value: 0.8594705677194221

[checkpoint_format]
ckpt_contains: bundle
bundle_keys: state_dict, rebuild_spec, log_summary

[common_cfg]
learning_rate: 0.00036594723903481
model_name: vade
encoder_type: recurrent
batch_size: 1024
latent_dim: 20
epochs: 50
n_components: 10
output_path: ./tutorial_files\tutorial_project\Trained_models
data_path: .
log_history: True
pretrained: None
save_weights: True
run: 0
num_workers: 0
prefetch_factor: 0
use_amp: False
interaction_regularization: 0.000348857818925854
kmeans_loss: 0.5
diag_max_batches: 4
seed: 0
limit_train_batches: 1000
limit_val_batches: 1000

[teacher_cfg]
use_turtle_teacher: True
teacher_gamma: 3
teacher_outer_steps: 500
teacher_inner_steps: 200
teacher_normalize_feats: True
teacher_head_temp: 0.35
teacher_task_temp: 0.35
teacher_alpha_sample_entropy: 1.0
lambda_distill: 8.0
lambda_decay_start: 8
lambda_end_weight: 0.0
lambda_cooldown: 12
distill_sharpen_T: 0.5
distill_conf_weight: True
distill_conf_thresh: 0.6
generic_lambda_distill: 2.0
generic_distill_sharpen_T: 0.5
generic_distill_conf_weight: True
generic_distill_conf_thresh: 0.6
generic_distill_warmup_epochs: 1
distill_class_reweight_beta: 1
distill_class_reweight_cap: 3
include_latent_view: True
include_edges_view: False
include_nodes_view: True
include_angles_view: False
pca_nodes_dim: 32
pca_edges_dim: 32
pca_angles_dim: 32
batch_size_nodes: 4096
batch_size_edges: 8192
batch_size_angles: 8192
teacher_refresh_every: None
teacher_freeze_at: 3
reinit_gmm_on_refresh: False
teacher_batch_size: 2048

[vade_cfg]
learning_rate_pretrain: 0.00183373647106535
gmm_learning_rate: 0.000122350019496003
pretrain_epochs: 10
reg_cat_clusters: 0.0
recluster: False
freeze_gmm_epochs: 0
freeze_decoder_epochs: 0
prior_loss_weight: 1.0
reg_scatter_weight: 0.0
temporal_cohesion_weight: 1.62686218712298
reg_scatter_beta: 0.0
repel_weight: 0.0
repel_length_scale: 0.0
tf_cluster_weight: 1.0
nonempty_weight: 0.0167301507709498
nonempty_p: 2.0
nonempty_floor_percent: 0.05
kmeans_loss_pretrain: 0.166630457111161
repel_weight_pretrain: 0.01
repel_length_scale_pretrain: 0.0
nonempty_weight_pretrain: 0.0
nonempty_p_pretrain: 2.0
nonempty_floor_percent_pretrain: 0.05
kl_annealing_mode: tf_sigmoid
kl_max_weight: 0.567943892625782
kl_warmup: 1
kl_end_weight: 1.0
kl_cooldown: 10
kl_annealing_mode_pretrain: tf_sigmoid
kl_max_weight_pretrain: 1.0
kl_warmup_pretrain: 1
kl_end_weight_pretrain: 1.0
kl_cooldown_pretrain: 10
========================================================================================

Let’S first can have a look at how the training went by inspecting the logging information within log_summary

deepof.visuals.plot_training_metrics(log_summary)

{'alignment': <Figure size 1000x460 with 1 Axes>,
 'loss': <Figure size 1000x480 with 1 Axes>}

NOTE: As we only use a single video for validation and a rather small training data set in these examples (which is why we only use a single vidoe for validation), we see some slightly odd differences between training and validation loss.

Here we can see that the validation loss followed the training loss quite nicely. The big gap here is explainable by variation between videos and the fact that we only used a single video for validation. The alignment score also quickly increased and remained strong during the entire training.

Now we can use our two model-checkpoints for inference, generating the latent, compressed represenation of our behavior-space (the embeddings) and the soft_counts (the percentage probability for each cluster according to the model)

torch.cuda.empty_cache()
# Get embeddings, soft_counts, and breaks per video
embeddings_val, soft_counts_val = deepof.clustering.model_utils_new.embedding_per_video(
    coordinates=my_deepof_project,
    to_preprocess=to_preprocess,
    model=model_val,
    meta_info=meta_info,
    animal_id="B",
    global_scaler=global_scaler,
)

Computing embeddings          : 100%|██████████| 53/53 [01:39<00:00,  1.87s/table]

torch.cuda.empty_cache()
# Get embeddings, soft_counts, and breaks per video
embeddings_score, soft_counts_score = deepof.clustering.model_utils_new.embedding_per_video(
    coordinates=my_deepof_project,
    to_preprocess=to_preprocess,
    model=model_score,
    meta_info=meta_info,
    animal_id="B",
    global_scaler=global_scaler,
)

Computing embeddings          : 100%|██████████| 53/53 [01:24<00:00,  1.60s/table]

Visualizing the embedding space reveals great anisotropy and a pretty clean separation between groups for both of our models

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_score,
    soft_counts_score,
    aggregate_experiments=None,
    samples=100,
    colour_by="cluster",
    ax=ax1,
    save=False,  # Set to True, or give a custom name, to save the plot
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_score,
    soft_counts_score,
    aggregate_experiments="time on cluster",  # Can also be set to 'mean' and 'median'
    exp_condition="CSDS", #"CSDS", SI, Treatment
    show_aggregated_density=False,
    ax=ax2,
    save=False,  # Set to True, or give a custom name, to save the plot,
)
ax2.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)

plt.tight_layout()
plt.show()

Info! Set colour_by to exp_condition as aggregate_experiments were given!

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_val,
    soft_counts_val,
    aggregate_experiments=None,
    samples=100,
    colour_by="cluster",
    ax=ax1,
    save=False,  # Set to True, or give a custom name, to save the plot
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_val,
    soft_counts_val,
    aggregate_experiments="time on cluster",  # Can also be set to 'mean' and 'median'
    exp_condition="CSDS", #"CSDS", SI, Treatment
    show_aggregated_density=False,
    ax=ax2,
    save=False,  # Set to True, or give a custom name, to save the plot,
)
ax2.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)

plt.tight_layout()
plt.show()

Info! Set colour_by to exp_condition as aggregate_experiments were given!

Now we can also use our supervised behaviors to see where they cluster in the embedding space.

supervised_annotation = my_deepof_project.supervised_annotation()

data preprocessing            : 100%|██████████| 4/4 [00:47<00:00, 11.82s/step, step=Loading kinematics]
supervised annotations        : 100%|██████████| 53/53 [01:23<00:00,  1.58s/table, step=post processing]

We do a few different subplots with different supervised behaviors for annotation.

fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 10))

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_score,
    soft_counts_score,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["B_stat-active","B_stat-passive","B_moving","B_stat-lookaround"],
    ax=ax1,
    save=False,
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_score,
    soft_counts_score,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["W_stat-active","W_stat-passive","W_moving","W_stat-lookaround"],
    ax=ax2,
    save=False,
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_score,
    soft_counts_score,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["W_sniff-arena","W_climb-arena","B_sniff-arena","B_climb-arena"],
    ax=ax3,
    save=False,
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_score,
    soft_counts_score,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["B_W_nose2nose","B_W_nose2body","B_W_following","W_B_following"],
    ax=ax4,
    save=False,
)

plt.tight_layout()

As we can see the embeddings space for the score checkpoint already separates the different movement categories of the black mouse to some extent with stationary-passive beavior clustering more to the top left and moving more to the bottom right. This separation is not really observable for the white mosue, which is expected, as we excluded it from the training data. The more specific behaviors like B_sniff-arena are less clearly structured.

fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 10))

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_val,
    soft_counts_val,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["B_stat-active","B_stat-passive","B_moving","B_stat-lookaround"],
    ax=ax1,
    save=False,
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_val,
    soft_counts_val,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["W_stat-active","W_stat-passive","W_moving","W_stat-lookaround"],
    ax=ax2,
    save=False,
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_val,
    soft_counts_val,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["W_sniff-arena","W_climb-arena","B_sniff-arena","B_climb-arena"],
    ax=ax3,
    save=False,
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings_val,
    soft_counts_val,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["B_W_nose2nose","B_W_nose2body","B_W_following","W_B_following"],
    ax=ax4,
    save=False,
)

plt.tight_layout()

the validation checkpoint shows similar results.

Another way to judge the quality of the embedding space more objectively (especially since the U-map representation can change to some degree based on the random state) is by using the plot_embedding_evaluation function. It plots three measures for a series of different (and selectable) behaviors. These measures are:

• Compactness, which denotes how “compact” the respective behavior within the embeddings space is e.g. wide or small scatter

• average precision, which denotes how well a simple LogisticRegression can predict specific behaviors given the embedding space

• k-nearest-neighbors, which denotes how well connected instances of the same behavior are in the embedding space

deepof.visuals.plot_embedding_evaluation(my_deepof_project,embeddings_score,supervised_annotation)

Evaluating embeddings:  62%|██████▎   | 10/16 [01:21<00:47,  7.89s/behavior]
Warning! Not enough instances found of behavior B_climb-arena within supervised_annotations.
Found 154, needed 200. You can adjust "minimum_number_of_positives" but measures may become imprecise.
Evaluating embeddings: 100%|██████████| 16/16 [01:49<00:00,  6.85s/behavior]

deepof.visuals.plot_embedding_evaluation(my_deepof_project,embeddings_val,supervised_annotation)

Evaluating embeddings:  62%|██████▎   | 10/16 [01:13<00:46,  7.83s/behavior]
Warning! Not enough instances found of behavior B_climb-arena within supervised_annotations.
Found 154, needed 200. You can adjust "minimum_number_of_positives" but measures may become imprecise.
Evaluating embeddings: 100%|██████████| 16/16 [01:42<00:00,  6.39s/behavior]

Here we see that, as expected, the simpler behaviors for the black mouse are also mroe directly encoded in the embeddings space and that the embeddings reflect behaviors of the White mosue and interactive behaviors much less, which is also expected based on our training data.

Let’s now investigate how well our model separates our stressed and non-stressed croup in the different clusters. Keep in mind that we did not provide any information regarding which animal is in which group at all. All separation here only arises from differences in the movement patterns of stressed and non-stressed animals.

fig, (ax, ax2) = plt.subplots(2, 1, figsize=(12, 5))

# over entire data
deepof.visuals.plot_enrichment(
    my_deepof_project,
    embeddings_score,
    soft_counts_score,
    normalize=False,
    add_stats="Mann-Whitney",
    exp_condition="CSDS",
    verbose=False,
    ax=ax,
)

# over only the first minute
deepof.visuals.plot_enrichment(
    my_deepof_project,
    embeddings_score,
    soft_counts_score,
    normalize=False,
    bin_size=60,
    bin_index=0,
    add_stats="Mann-Whitney",
    exp_condition="CSDS",
    verbose=False,
    ax=ax2,
)
ax.tick_params(
    axis='x',
    which='both',
    bottom=False,
    top=False,
    labelbottom=False)
ax.set_xlabel("")
ax2.legend().remove()
plt.title("")
plt.tight_layout()
plt.show()

# The same for the validation checkpoint

fig, (ax, ax2) = plt.subplots(2, 1, figsize=(12, 5))

deepof.visuals.plot_enrichment(
    my_deepof_project,
    embeddings_val,
    soft_counts_val,
    normalize=False,
    add_stats="Mann-Whitney",
    exp_condition="CSDS",
    verbose=False,
    ax=ax,
)

deepof.visuals.plot_enrichment(
    my_deepof_project,
    embeddings_val,
    soft_counts_val,
    normalize=False,
    bin_size=60,
    bin_index=0,
    add_stats="Mann-Whitney",
    exp_condition="CSDS",
    verbose=False,
    ax=ax2,
)
ax.tick_params(
    axis='x',
    which='both',
    bottom=False,
    top=False,
    labelbottom=False)
ax.set_xlabel("")
ax2.legend().remove()
plt.title("")
plt.tight_layout()
plt.show()

We see some very clear differences here, which is expected, as both groups are, indeed, pretty different. In terms of how to proceed with the analysis we now refer to the unsupervised tutorial. Here we will be content with only extracting some video snippets of our clusters.

Looking at cluster 0, that is strongly populated by controls and cluster 1, for which the reverse is true, we immediately see a clear difference

# uncomment this to export the videos the snippets above came from

'''
deepof.visuals.export_annotated_video(
    my_deepof_project,
    soft_counts = soft_counts_val,
    supervised_annotations = None,
    # Time selection parameters
    frame_limit_per_video = 100,
    #others
    behaviors = ["Cluster_0","Cluster_1"],
    experiment_id = None,
    min_confidence = 0.75,
    min_bout_duration = None,
    display_time = True,
    display_counter = True,
    display_arena = True,
    display_markers = True,
    #display_mouse_labels = True,
    cluster_names = None,
)
'''

'\ndeepof.visuals.export_annotated_video(\n    my_deepof_project,\n    soft_counts = soft_counts_val,\n    supervised_annotations = None,\n    # Time selection parameters\n    frame_limit_per_video = 100,\n    #others\n    behaviors = ["Cluster_0","Cluster_1"],\n    experiment_id = None,\n    min_confidence = 0.75,\n    min_bout_duration = None,\n    display_time = True,\n    display_counter = True,\n    display_arena = True,\n    display_markers = True,\n    #display_mouse_labels = True,\n    cluster_names = None,\n)\n'

Whilst VaDE is good at identifying single mouse patterns, it often struggles with multiple mice, among other reasons because of its sensibility to scale changes. Respectively, for multi-mouse embeddings we can use a contrastive model which does not have this weakness.

How to train a Contrastive model for paired mouse behaviors

The steps in preparing the data, training the model, and evaluation mostly essentially the same. The first noteworthy difference is that we do not state an animal id during data preprocessing. This wil lresult in all of our animals being used for training including the distance information between them (and longer training times)

# This code will generate a dataset using graph representations, as well a some auxiliary objects
preprocessed_object, meta_info, adj_matrix, to_preprocess, global_scaler = my_deepof_project.get_graph_dataset(
    center="Center",
    align="Spine_1",
    window_size=24,
    window_step=1,
    test_videos=1,
    preprocess=True,
    scale="standard",
    dist_standardize = "groupwise",
    speed_standardize = "groupwise",
    coord_standardize = "groupwise",
)
samples_max = None
bin_size = None
bin_index = None
precomputed_bins = None

Loading tables                : 100%|██████████| 5/5 [00:21<00:00,  4.34s/step, step=Get graph info]
Sampling                      : 100%|██████████| 53/53 [00:05<00:00,  9.28table/s]
Scaling                       : 100%|██████████| 53/53 [00:22<00:00,  2.36table/s]
Get training windows          : 100%|██████████| 52/52 [00:00<00:00, 298.86table/s]
Get testing windows           : 100%|██████████| 1/1 [00:00<00:00, 63.48table/s]
Reshaping                     : 100%|██████████| 2/2 [00:05<00:00,  2.66s/table]

import deepof.clustering.training
model_val, model_score, model_part, log_summary = deepof.clustering.training.train_deepof_model(


     # As before, we load a pretrained model. Setting this path will lead to a model getting loaded and all other inputs getting ignored.
    pretrained="./tutorial_files/tutorial_project/Trained_models/models/contrastive/run_0/best_model_val.pth",

    # Basics
    output_path = os.path.join(my_deepof_project._project_path, my_deepof_project._project_name, "Trained_models"), #where to save the model after training
    preprocessed_object=preprocessed_object, # inputs from your preprocessing
    adjacency_matrix=adj_matrix,
    meta_info=meta_info,
    model_name="Contrastive",       # Determines the model-type, can be VaDE, VQVAE or Contrastive
    encoder_type="TCN",             # Determines the type of "funnel" to use to compress your data into the latent space. Can be recurrent, TCN or transformer
    batch_size=256,
    latent_dim=20,                  # Determins the size of the latent/embedding space
    epochs=50,
    run=0,                          # run id, change this to change the name of the folder your model will be saved in
    learning_rate=0.00025,          # learning rates for the Contrastive model

    # Losses
    contrastive_similarity_function = "cosine",
    contrastive_loss_function = "nce",

    # Augmentations
    aug_min_shift = 1, # Shifts entire frame window randomly forwards and backwards in time
    aug_max_shift = 4,
    aug_p_shift = 0.6,

    aug_max_rot = 30,  # Adds random rotations to mouse bodyparts
    aug_n_rot = 3,
    aug_p_rot = 0.0,

    aug_max_interp = 8, # Interpolates random segments of mouse coordinates
    aug_min_interp = 3,
    aug_p_interp = 0.6,

    aug_noise_sigma = 0.03, # adds random x- or y-axis noise
    aug_p_noise = 0.0,

    # Teacher model setup
    use_turtle_teacher=False, # we do not use the Turtle teacher here

)

Using device: cuda
BatchDictDataset: building HDF5 at ./tutorial_files\tutorial_project\Trained_models\Datasets...
  Reason: X shape mismatch
HDF5 built. Samples: 778120, chunks: 256
HDF5 dataset ready. Samples: 778120, x_shape: (24, 22, 3), a_shape: (24, 26, 1), ang_shape: None
BatchDictDataset: building HDF5 at ./tutorial_files\tutorial_project\Trained_models\Datasets...
  Reason: X shape mismatch
HDF5 built. Samples: 14977, chunks: 256
HDF5 dataset ready. Samples: 14977, x_shape: (24, 22, 3), a_shape: (24, 26, 1), ang_shape: None
TensorBoard logs -> ./tutorial_files\tutorial_project\Trained_models\logs\contrastive_run_0
Loading pretrained weights from ./tutorial_files/tutorial_project/Trained_models/models/contrastive/run_11/best_model_val.pth

========================================================================================
Model info: tutorial_files\tutorial_project\Trained_models\models\contrastive\run_11\best_model_val_info.txt
----------------------------------------------------------------------------------------
stage: best_val
epoch: 46
train_steps: 142880
val_total: 2.4080547437829485

[checkpoint_format]
ckpt_contains: bundle
bundle_keys: state_dict, rebuild_spec, log_summary

[common_cfg]
learning_rate: 0.00025
model_name: contrastive
encoder_type: TCN
batch_size: 256
latent_dim: 20
epochs: 50
n_components: 10
output_path: ./tutorial_files\tutorial_project\Trained_models
data_path: .
log_history: True
pretrained: None
save_weights: True
run: 11
num_workers: 0
prefetch_factor: 0
use_amp: False
interaction_regularization: 0.0
kmeans_loss: 0.0
diag_max_batches: 4
seed: 0
limit_train_batches: 1000
limit_val_batches: 1000

[teacher_cfg]
use_turtle_teacher: False
teacher_gamma: 8.0
teacher_outer_steps: 500
teacher_inner_steps: 100
teacher_normalize_feats: True
teacher_head_temp: 0.5
teacher_task_temp: 0.5
teacher_alpha_sample_entropy: 2.0
lambda_distill: 4.0
lambda_decay_start: 10
lambda_end_weight: 0.2
lambda_cooldown: 10
distill_sharpen_T: 0.5
distill_conf_weight: False
distill_conf_thresh: 0.3
generic_lambda_distill: 2.0
generic_distill_sharpen_T: 0.5
generic_distill_conf_weight: True
generic_distill_conf_thresh: 0.6
generic_distill_warmup_epochs: 1
distill_class_reweight_beta: 1
distill_class_reweight_cap: 3
include_latent_view: False
include_edges_view: False
include_nodes_view: True
include_angles_view: False
pca_nodes_dim: 32
pca_edges_dim: 32
pca_angles_dim: 32
batch_size_nodes: 4096
batch_size_edges: 8192
batch_size_angles: 8192
teacher_refresh_every: None
teacher_freeze_at: 10
reinit_gmm_on_refresh: False
teacher_batch_size: 2048

[contrastive_cfg]
temperature: 0.1
contrastive_similarity_function: cosine
contrastive_loss_function: nce
beta: 0.1
tau: 0.1
aug_min_shift: 1
aug_max_shift: 4
aug_p_shift: 0.6
aug_max_rot: 30
aug_n_rot: 3
aug_p_rot: 0.0
aug_max_interp: 8
aug_min_interp: 3
aug_p_interp: 0.6
aug_noise_sigma: 0.03
aug_p_noise: 0.0
========================================================================================

Whilst the command is exactly the same, the metrics for the contrastive model differ slightly, as instead of the score, positive and negative similarities are shown

deepof.visuals.plot_training_metrics(log_summary)

{'similarities': <Figure size 1000x480 with 1 Axes>,
 'loss': <Figure size 1000x480 with 1 Axes>}

NOTE: As we only use a single video for validation and a rather small training data set in these examples (which is why we only use a single vidoe for validation), we see some slightly odd differences between training and validation loss.

The first major difference between the Contrastive model and Vade becomes visible at inference. Because the contrastive model has no decoder, we need to decode the embeddings manually. This can be done by choosing an extraction method which can be:

• “gmm” fitting gaussian mixture models to the embedding space

• “msm” fitting a msm-pcca model which identifies strongly connected subsectiosn oft the embedding space

• “combined” fits msm first, then gmm to sections of the embedding space, that correspond to frames with low tracking accuracy

Then the embedding_gates have to be chosen. These can either be a bodypart, which will result in the embeddings being split based on mouse-mouse distances between this bodypart into a “mice close”, “mice medium distance” and “mice far away” subsection which are clustered seperately. Alternatively embedding_gates also accepts behaviors that are then used to split the embeddings into “behavior present” and “behavior absent”. This functionality allows us to pre-organize the clustering for better results.

Finally states_per_gate denotes the number of states per gate. for the current configuration we will get 8*3 (3 embedding_gates) + 8 (GMM_gate from combined mode) = 24 clusters

import deepof.clustering.model_utils_new
torch.cuda.empty_cache()
# Get embeddings, soft_counts, and breaks per video
embeddings, soft_counts = deepof.clustering.model_utils_new.embedding_per_video(
    coordinates=my_deepof_project,
    to_preprocess=to_preprocess,
    model=model_val,
    meta_info=meta_info,
    extract_pair=["B","W"],
    softcounts_extraction_method="combined",
    embedding_gates = "Center",
    states_per_gate=8,
    global_scaler=global_scaler,
)

Computing embeddings          : 100%|██████████| 53/53 [00:59<00:00,  1.11s/table]
Fit MSM/PCCA                  : 100%|██████████| 3/3 [00:15<00:00,  5.20s/model]
Decode softcounts (MSM/PCCA)  : 100%|██████████| 53/53 [00:03<00:00, 14.05table/s]
Fit GMMs                      : 100%|██████████| 2/2 [00:19<00:00,  9.63s/gmm]
Decode softcounts (GMM)       : 100%|██████████| 53/53 [00:01<00:00, 45.34table/s]

Now we proceed with the same visualizations as for the VaDE Model before.

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings,
    soft_counts,
    aggregate_experiments=None,
    samples=100,
    colour_by="cluster",
    ax=ax1,
    save=False,  # Set to True, or give a custom name, to save the plot
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings,
    soft_counts,
    aggregate_experiments="time on cluster",  # Can also be set to 'mean' and 'median'
    exp_condition="CSDS", #"CSDS", SI, Treatment
    show_aggregated_density=False,
    ax=ax2,
    save=False,  # Set to True, or give a custom name, to save the plot,
)
ax2.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.0)

plt.tight_layout()
plt.show()

Info! Set colour_by to exp_condition as aggregate_experiments were given!

fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 10))

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings,
    soft_counts,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["B_stat-active","B_stat-passive","B_moving","B_stat-lookaround"],
    ax=ax1,
    save=False,
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings,
    soft_counts,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["W_stat-active","W_stat-passive","W_moving","W_stat-lookaround"],
    ax=ax2,
    save=False,
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings,
    soft_counts,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["W_sniff-arena","W_climb-arena","B_sniff-arena","B_climb-arena"],
    ax=ax3,
    save=False,
)

deepof.visuals.plot_embeddings(
    my_deepof_project,
    embeddings,
    soft_counts,
    supervised_annotation,
    aggregate_experiments=None,
    samples=100,
    colour_by=["B_W_nose2nose","B_W_nose2body","B_W_following","W_B_following"],
    ax=ax4,
    save=False,
)

plt.tight_layout()

Here we now can see how the social behaviors cluster considerably better in the embedding space. We can also see that the model recognizes directionality by following being separated by which mouse follows which.

deepof.visuals.plot_embedding_evaluation(my_deepof_project,embeddings,supervised_annotation)

Evaluating embeddings:  62%|██████▎   | 10/16 [01:12<00:38,  6.39s/behavior]
Warning! Not enough instances found of behavior B_climb-arena within supervised_annotations.
Found 98, needed 200. You can adjust "minimum_number_of_positives" but measures may become imprecise.
Evaluating embeddings: 100%|██████████| 16/16 [01:25<00:00,  5.34s/behavior]

This observation is also confirmed by the higher preformance measures for social behaviors

fig, (ax, ax2) = plt.subplots(2, 1, figsize=(12, 5))

# over entire data
deepof.visuals.plot_enrichment(
    my_deepof_project,
    embeddings,
    soft_counts,
    normalize=False,
    add_stats="Mann-Whitney",
    exp_condition="CSDS",
    verbose=False,
    ax=ax,
)

# over only the first minute
deepof.visuals.plot_enrichment(
    my_deepof_project,
    embeddings,
    soft_counts,
    normalize=False,
    bin_size=60,
    bin_index=0,
    add_stats="Mann-Whitney",
    exp_condition="CSDS",
    verbose=False,
    ax=ax2,
)
ax.tick_params(
    axis='x',
    which='both',
    bottom=False,
    top=False,
    labelbottom=False)
ax.set_xlabel("")
ax2.legend().remove()
plt.title("")
plt.tight_layout()
plt.show()

We again see clear separation between stressed and non-stressed mice. even within the much smaller last 8 GMM clusters some separation is observable, even though we expect that these clusters are much less consitent in general.

Let’s have a look at some clusters, specifically cluster 1 (close mice, ms-pcca) and cluster 31 (uncertain frames, GMM)

here we can see that cluster 1, as well as many msm-pcca-based clusters, are often highly consistent in terms of the relative orientation and interaction of the mice with each other

In contrast, the gmm clusters capture less consistent but also interesting behavior (in this cluster specifically we see a lot of fighting but also other behaviors)

# uncomment this to export the videos the snippets above came from

'''
deepof.visuals.export_annotated_video(
    my_deepof_project,
    soft_counts = soft_counts,
    supervised_annotations = None,
    # Time selection parameters
    frame_limit_per_video = 100,
    #others
    #behaviors = ["Cluster_1","Cluster_26"],
    experiment_id = None,
    min_confidence = 0.1, # as we have many more clusters and as the contrastive model produces low confidence spearations, we reduce the minimum confidence level
    min_bout_duration = None,
    display_time = True,
    display_counter = True,
    display_arena = True,
    display_markers = True,
    #display_mouse_labels = True,
    cluster_names = None,
)
'''

Exporting behavior videos     : 100%|██████████| 32/32 [04:37<00:00,  8.67s/ video]

What’s next

With this we conclude the model training tutorial. Next, we can also explore “the other branch” of the package: we’ll learn to use deepOFs wide array of pre-defined behaviors or even completely user-defined behaviors to measure exactly what we are interested in within our data. In short the supervised branch for detecting pre-defined behaviors!