Welcome to BehaveNet’s documentation!

BehaveNet is a probabilistic framework for the analysis of behavioral video and neural activity. This framework provides tools for compression, segmentation, generation, and decoding of behavioral videos. Please see the original paper and the follow-up paper on semi-supervised autoencoders for additional details.

Installation

Before you begin, ensure that you have downloaded both Git (for cloning the repository) and Anaconda (for managing the code environment). The following steps will show you how to:

  1. Set up a virtual environment using Anaconda

  2. Install the BehaveNet and ssm packages

  3. Store user paths in a local json file

Environment setup

First set up the anaconda environment from the command line.

$: conda create --name=behavenet python=3.7.2
$: conda activate behavenet

Package installation

BehaveNet

cd to the directory that will contain the BehaveNet code, and then clone the BehaveNet repository from github:

(behavenet) $: git clone https://github.com/ebatty/behavenet

cd into the behavenet repository and install the required dependencies:

(behavenet) $: pip install -r requirements.txt

To make the package modules visible to the python interpreter, locally run pip install from inside the main behavenet directory:

(behavenet) $: pip install -e .

You can test the installation by running

(behavenet) $: python -c "import behavenet"

If this command does not return an error the package has been successfully installed.

Installing the BehaveNet package automatically installed the ipykernel package, which allows you to work with python code in Jupyter notebooks. To be able to use the behavenet conda environment for Jupyter notebooks, run the following command from the terminal:

(behavenet) $: python -m ipykernel install --user --name behavenet

ssm

The ssm package is the backend state space modeling code used by BehaveNet. To install ssm, cd to any directory where you would like to keep the ssm code and run the following:

(behavenet) $: git clone https://github.com/slinderman/ssm.git --branch behavenet-no-cython --single-branch
(behavenet) $: cd ssm
(behavenet) $: pip install cython
(behavenet) $: pip install -e .

PyTorch

PyTorch is automatically pip installed during the BehaveNet installation process; however, if you have issues running PyTorch, first uninstall the existing package:

(behavenet) $: pip uninstall torch

and reinstall following the directions here using the Pip package option.

ffmpeg

The BehaveNet package uses the ffmpeg backend to produce movies. ffmpeg is automatically installed on many systems, and is not automatically installed with BehaveNet. If you are trying to make movies and run into issues with ffmpeg, install using the conda package manager:

(behavenet) $: conda install -c conda-forge ffmpeg

Set user paths

Next, set up your paths to the directories where data, results, and figures will be stored. To do so, launch python from the behavenet environment, and type:

from behavenet import setup
setup()

You will be asked to input a base data directory; all data should be stored in the form base_data_dir/lab_id/expt_id/animal_id/session_id/data.hdf5. More information on the structure of the hdf5 file can be found here. You will also be asked to input a base results directory, which will store all of the model fits. Finally, the base figure directory will be used to store figure and video outputs.

The behavenet.setup() function will create a hidden directory named .behavenet in your user directory.

  • In Linux, ~/.behavenet

  • In MacOS, /Users/CurrentUser/.behavenet

Within this directory the function will create a json file named directories which you can manually edit at any point.

User guide

Introduction

BehaveNet is a software package that provides tools for analyzing behavioral video and neural activity. Currently BehaveNet supports:

  • Video compression using convolutional autoencoders

  • Video segmentation (and generation) using autoregressive hidden Markov models

  • Neural network decoding of videos from neural activity

  • Bayesian decoding of videos from neural activity

BehaveNet automatically saves models using a well-defined and flexible directory structure, allowing for easy management of many models and multiple datasets.

The command line interface

Users interact with BehaveNet using a command line interface, so all model fitting is done from the terminal. To simplify this process all necessary parameters are defined in four configuration files that can be manually updated using a text editor:

  • data_config - dataset ids, video frames sizes, etc. You can automatically generate this configuration file for a new dataset by following the instructions in the following section.

  • model_config - model hyperparameters

  • training_config - learning rate, number of epochs, etc.

  • compute_config - gpu vs cpu, gpu ids, etc.

Example configuration files can be found here.

For example, the command line call to fit an autoencoder would be (using the default json files):

$: cd /behavenet/code/directory/
$: cd behavenet
$: python fitting/ae_grid_search.py --data_config ../configs/data_default.json --model_config ../configs/ae_jsons/ae_model.json --training_config ../configs/ae_jsons/ae_training.json --compute_config ../configs/ae_jsons/ae_compute.json

We recommend that you copy the default config files in the behavenet repo into a separate directory on your local machine and make edits there. For more information on the different hyperparameters, see the hyperparameters glossary.

Adding a new dataset

When using BehaveNet with a new dataset you will need to make a new data config json file, which can be automatically generated using a BehaveNet helper function. You will be asked to enter the following information (examples shown for Musall dataset):

  • lab or experimenter name (musall)

  • experiment name (vistrained)

  • example animal name (mSM36)

  • example session name (05-Dec-2017)

  • input channels (2) - this can refer to color channels (for RGB data) and/or number of camera views, which should be concatenated along the color channel dimension. In the Musall dataset we use grayscale images from two camera views, so a trial with 189 frames will have a block of video data of shape (189, 2, 128, 128)

  • y pixels (128)

  • x pixels (128)

  • use output mask (False) - an optional output mask can be applied to each video frame if desired; these output masks must also be stored in the data.hdf5 files as masks.

  • frame rate (30) - in Hz; BehaveNet assumes that the video data and neural data are binned at the same temporal resolution

  • neural data type (ca) - either ca for 2-photon/widefield data, or spikes for ephys data. This parameter controls the noise distribution for encoding models, as well as several other model hyperparameters.

To enter this information, launch python from the behavenet environment and type:

from behavenet import add_dataset
add_dataset()

This function will create a json file named [lab_id]_[expt_id].json in the .behavenet directory in your user home directory, which you can manually update at any point using a text editor.

Organizing model fits with test-tube

BehaveNet uses the test-tube package to organize model fits into user-defined experiments, log meta and training data, and perform grid searches over model hyperparameters. Most of this occurs behind the scenes, but there are a couple of important pieces of information that will improve your model fitting experience.

BehaveNet organizes model fits using a combination of hyperparameters and user-defined experiment names. For example, let’s say you want to fit 5 different convolutional autoencoder architectures, all with 12 latents, to find the best one. Let’s call this experiment “arch_search”, which you will set in the model_config json in the experiment_name field. The results will then be stored in the directory results_dir/lab_id/expt_id/animal_id/session_id/ae/conv/12_latents/arch_search/.

Each model will automatically be assigned it’s own “version” by test-tube, so the arch_search directory will have subdirectories version_0, …, version_4. If an additional CAE model is later fit with 12 latents (and using the “arch_search” experiment name), test-tube will add it to the arch_search directory as version_5. Different versions may have different architectures, learning rates, regularization values, etc. Each model class (autoencoder, arhmm, decoders) has a set of hyperparameters that are used for directory names, and another set that are used to distinguish test-tube versions within the user-defined experiment.

Within the version_x directory, there are various files saved during training. Here are some of the files automatically output when training an autoencoder:

  • best_val_model.pt: the best model (not necessarily from the final training epoch) as determined by computing the loss on validation data

  • meta_tags.csv: hyperparameters associated with data, computational resources, training, and model

  • metrics.csv: metrics computed on dataset as a function of epochs; the default is that metrics are computed on training and validation data every epoch (and reported as a mean over all batches) while metrics are computed on test data only at the end of training using the best model (and reported per batch).

  • session_info.csv: experimental sessions used to fit the model

Additionally, if you set export_latents to True in the training config file, you will see

  • [lab_id]_[expt_id]_[animal_id]_[session_id]_latents.pkl: list of np.ndarrays of CAE latents computed using the best model

and if you set export_train_plots to True in the training config file, you will see

  • loss_training.png: plot of MSE as a function of training epoch on training data

  • loss_validation.png: same as above using validation data

Grid searching with test-tube

Beyond organizing model fits, test-tube is also useful for performing grid searches over model hyperparameters, using multiple cpus or gpus. All you as the user need to do is enter the relevant hyperparameter choices as a list instead of a single value in the associated configuration file.

Again using the autoencoder as an example, let’s say you want to fit a single AE architecture using 4 different numbers of latents, all with the same regularization value. In the model config file, you will set these values as:

{
...
"n_ae_latents": [4, 8, 12, 16],
"l2_reg": 0.0,
...
}

To specify the computing resources for this job, you will next edit the compute config file, which looks like this:

{
...
"device": "cuda", # "cpu" or "cuda"
"gpus_viz": "0", # "add multiple gpus as e.g. "0;1;4"
"tt_n_gpu_trials": 1000,
"tt_n_cpu_trials": 1000,
"tt_n_cpu_workers": 5,
...
}

With the device field set to cuda, test-tube will use gpus to run this job. The gpus_viz field can further specify which subset of gpus to use. The tt_n_gpu_trials defines the maximum number of jobs to run. If this number is larger than the total number of hyperparameter configurations, all configurations are fit; if this number is smaller than the total number (say if "tt_n_gpu_trials": 2 in this example) then this number of configurations is randomly sampled from all possible choices.

To fit models using the cpu instead, set the device field to cpu; then tt_n_cpu_workers defines the total number of cpus to run the job (total number of models fitting at any one time) and tt_n_cpu_trials is analogous to tt_n_gpu_trials.

Finally, multiple hyperparameters can be searched over simultaneously; for example, to search over both AE latents and regularization values, set these parameters in the model config file like so:

{
...
"n_ae_latents": [4, 8, 12, 16],
"l2_reg": [1e-5, 1e-4, 1e-3],
...
}

This job would then fit a total of 4 latent values x 3 regularization values = 12 models.

Autoencoders

BehaveNet uses convolutional autoencoders to perform nonlinear dimensionality reduction on behavioral videos. The steps below demonstrate how to fit these models on an arbitrary dataset.

Within the behavenet package there is a directory named configs, which contains example config files. First copy the following config files to the .behavenet directory that was automatically created in your home directory: ae_compute.json, ae_model.json, and ae_training.json. You can then update the hyperparameters in these files in a text editor.

To fit a single model with the default CAE BehaveNet architecture (details in paper), edit the ae_model.json file to look like the following:

{
"experiment_name": "ae-example",
"model_type": "conv",
"n_ae_latents": 12,
"l2_reg": 0.0,
"rng_seed_model": 0,
"fit_sess_io_layers": false,
"ae_arch_json": null,
"model_class": "ae"
}

Then to fit the model, cd to the behavenet directory in the terminal and run

$: python behavenet/fitting/ae_grid_search.py --data_config /user_home/.behavenet/musall_vistrained_params.json --model_config /user_home/.behavenet/ae_model.json --training_config /user_home/.behavenet/ae_training.json --compute_config /user_home/.behavenet/ae_compute.json

where ~/.behavenet/musall_vistrained_params.json can be replaced by any dataset config file created by running the behavenet.add_dataset() function (example here).

Performing a search over multiple latents is as simple as editing the ae_model.json as below and rerunning the same command.

{
"experiment_name": "latent-search",
"model_type": "conv",
"n_ae_latents": [6, 9, 12, 15],
"l2_reg": 0.0,
"rng_seed_model": 0,
"fit_sess_io_layers": false,
"ae_arch_json": null,
"model_class": "ae"
}

Training an AE can be slow: you can speed up the training by parallelizing over multiple gpus. To do this, just specify n_parallel_gpus to be the number of gpus you wish to use per model. The code will split up the gpus specified in gpus_viz into groups of size n_parallel_gpus (or less if there are leftover gpus) and run the models accordingly.

Conditional autoencoders

One drawback to the use of unsupervised dimensionality reduction (performed by the convolutional autoencoder) is that the resulting latents are generally uninterpretable, because any animal movement in a behavioral video will be represented across many (if not all) of the latents. Thus there is no simple way to find an “arm” dimension that is separate from a “pupil” dimension, distinctions that may be important for downstream analyses.

Semi-supervised approaches to dimensionality reduction offer a partial resolution to this problem. In this framework, the user first collects a set of markers that track body parts of interest over time. These markers can be, for example, the output of standard pose estimation software such as DeepLabCut, LEAP, or DeepPoseKit. These markers can then be used to augment the latent space (using conditional autoencoders) or regularize the latent space (using the matrix subspace projection loss), both of which are described below.

In order to fit these models, the data HDF5 needs to be augmented to include a new HDF5 group named labels, which contains an HDF5 dataset for each trial. The labels for each trial must match up with the corresponding video frames; for example, if the image data in images/trial_0013 contains 100 frames (a numpy array of shape [100, n_channels, y_pix, x_pix]), the label data in labels/trial_0013 should contain the corresponding labels (a numpy array of shape [100, n_labels]). See the data structure documentation for more information).

Conditional autoencoders

The conditional autoencoder implemented in BehaveNet is a simple extension of the convolutional autoencoder. Each frame is pushed through the encoder to produce a set of latents, which are concatenated with the corresponding labels; this augmented vector is then used as input to the decoder.

To fit a single conditional autoencoder with the default CAE BehaveNet architecture, edit the model_class parameter of the ae_model.json file:

{
"experiment_name": "ae-example",
"model_type": "conv",
"n_ae_latents": 12,
"l2_reg": 0.0,
"rng_seed_model": 0,
"fit_sess_io_layers": false,
"ae_arch_json": null,
"model_class": "cond-ae",
"conditional_encoder": false
}

Then to fit the model, use the ae_grid_search.py function using this updated model json. All other input jsons remain unchanged.

By concatenating the labels to the latents, we are learning a conditional decoder. We can also condition the latents on the labels by learning a conditional encoder. Turning on this feature requires an additional HDF5 group; documentation coming soon.

Matrix subspace projection loss

An alternative way to obtain a more interpretable latent space is to encourage a subspace to predict the labels themselves, rather than appending them to the latents. With appropriate additions to the loss function, we can ensure that the subspace spanned by the label-predicting latents is orthogonal to the subspace spanned by the remaining unconstrained latents. This is the idea of the matrix subspace projection loss.

For example, imagine we are tracking 4 body parts, each with their own x-y coordinates for each frame. This gives us 8 dimensions of behavior to predict. If we fit a CAE with 10 latent dimensions, we will use 8 of those dimensions to predict the 8 marker dimensions - one latent dimension for each marker dimension. This leaves 2 unconstrained dimensions to predict remaining variability in the images not captured by the labels. The model is trained by minimizing the mean square error between the true and predicted images, as well as the true and predicted labels. Unlike the conditional autoencoder described above, this new loss function has an additional hyperparameter that governs the tradeoff between image reconstruction and label reconstruction.

To fit a single autoencoder with the matrix subspace projection loss (and the default CAE BehaveNet architecture), edit the model_class and msp.alpha parameters of the ae_model.json file:

{
"experiment_name": "ae-example",
"model_type": "conv",
"n_ae_latents": 12,
"l2_reg": 0.0,
"rng_seed_model": 0,
"fit_sess_io_layers": false,
"ae_arch_json": null,
"model_class": "cond-ae-msp",
"msp.alpha": 1e-4,
"conditional_encoder": false
}

The msp.alpha parameter needs to be tuned for each dataset, but msp.alpha=1.0 is a reasonable starting value if the labels have each been z-scored.

Note

The matrix subspace projection model implemented in BehaveNet learns a linear mapping from the original latent space to the predicted labels that does not contain a bias term. Therefore you should center each label before adding them to the HDF5 file. Additionally, normalizing each label by its standard deviation can make searching across msp weights less dependent on the size of the input image.

Then to fit the model, use the ae_grid_search.py function using this updated model json. All other input jsons remain unchanged.

Partitioned subspace variational autoencoder

One downside to the MSP model introduced in the previous section is that the representation in the unsupervised latent space may be difficult to interpret. The partitioned subspace VAE (PS-VAE) attempts to remedy this situation by encouraging the unsupervised representation to be factorized, which has shown to help with interpretability (see paper here).

To fit a single PS-VAE (and the default CAE BehaveNet architecture), edit the model_class, ps_vae.alpha, ps_vae.beta, and ps_vae.anneal_epochs parameters of the ae_model.json file:

{
"experiment_name": "ae-example",
"model_type": "conv",
"n_ae_latents": 12,
"l2_reg": 0.0,
"rng_seed_model": 0,
"fit_sess_io_layers": false,
"ae_arch_json": null,
"model_class": "ps-vae",
"ps_vae.alpha": 1000,
"ps_vae.beta": 10,
"ps_vae.anneal_epochs": 100,
"conditional_encoder": false
}

The ps_vae.alpha and ps_vae.beta parameters need to be tuned for each dataset. See the guidelines for setting these parameters here.

Then to fit the model, use the ae_grid_search.py function using this updated model json. All other input jsons remain unchanged. See the hyperparameter search guide for information on how to efficiently search over the ps_vae.alpha and ps_vae.beta hyperparameters.

Multi-session PS-VAE

The Partitioned Subspace VAE (PS-VAE) (see preprint here) finds a low-dimensional latent representation of a single behavioral video that is partitioned into two subspaces: a supervised subspace that reconstructs user-provided labels, and an unsupervised subspace that captures remaining variability. In practice, though, we will typically want to produce a low-dimensional latent representation that is shared across multiple experimental sessions, rather than fitting session-specific models. However, the inclusion of multiple videos during training introduces a new problem: different videos from the same experimental setup will contain variability in the experimental equipment, lighting, or even physical differences between animals, despite efforts to standardize these features. We do not want these differences (which we refer to collectively as the “background”) to contaminate the latent representation, as they do not contain the behavioral information we wish to extract for downstream analyses.

To address this issue within the framework of the PS-VAE, we introduce a new subspace into our model which captures static differences between sessions (the “background” subspace) while leaving the other subspaces (supervised and unsupervised) to capture dynamic behaviors.

As with the PS-VAE, the data HDF5 needs to be augmented to include a new HDF5 group named labels, which contains an HDF5 dataset for each trial. The labels for each trial must match up with the corresponding video frames; for example, if the image data in images/trial_0013 contains 100 frames (a numpy array of shape [100, n_channels, y_pix, x_pix]), the label data in labels/trial_0013 should contain the corresponding labels (a numpy array of shape [100, n_labels]). See the data structure documentation for more information. Also see the documentation for fitting models on multiple sessions for more information on how to specify which sessions are used for fitting in the data json.

To fit an MSPS-VAE with the default CAE BehaveNet architecture, edit the model_class, ps_vae.alpha, ps_vae.beta, ps_vae.delta, n_background, n_sessions_per_batch and ps_vae.anneal_epochs parameters of the ae_model.json file:

{
"experiment_name": "ae-example",
"model_type": "conv",
"n_ae_latents": 12,
"l2_reg": 0.0,
"rng_seed_model": 0,
"fit_sess_io_layers": false,
"ae_arch_json": null,
"model_class": "msps-vae",
"ps_vae.alpha": 1000,
"ps_vae.beta": 10,
"ps_vae.delta": 50,
"ps_vae.anneal_epochs": 100,
"n_background": 3,
"n_sessions_per_batch": 2,
"conditional_encoder": false
}

The n_background parameter sets the dimensionality of the background subspace; we find 3 works well in practice. The n_sessions_per_batch parameter determines how many many sessions comprise a single batch during training; this value should be greater than 1 for the triplet loss to work. The current implementation supports values of n_sessions_per_batch = [2, 3, 4].

To fit the model, use the ae_grid_search.py function using this updated model json. You will also need to update the data json as detailed here. See the hyperparameter search guide for information on how to efficiently search over the ps_vae.alpha, ps_vae.beta, and ps_vae.delta hyperparameters.

ARHMMs

The next step of the BehaveNet pipeline is to model the low-dimensional representation of behavior with a simple class of nonlinear dynamical systems called autoregressive hidden Markov models (ARHMMs). An ARHMM models the sequence of continuous latents as a stochastic process that switches between a small number K of discrete states, each characterized by linear-Gaussian dynamics. These discrete state variables also exhibit temporal dependences through Markovian dynamics - the discrete state at time t may depend on its preceding value.

Fitting a single ARHMM is very similar to the AE fitting procedure; first copy the example json files arhmm_compute.json, arhmm_model.json, and arhmm_training.json into your .behavenet directory, cd to the behavenet directory in the terminal, and run:

$: python behavenet/fitting/arhmm_grid_search.py --data_config /user_home/.behavenet/musall_vistrained_params.json --model_config /user_home/.behavenet/arhmm_model.json --training_config /user_home/.behavenet/arhmm_training.json --compute_config /user_home/.behavenet/arhmm_compute.json

Decoders

The next step of the BehaveNet pipeline uses the neural activity to decode (or reconstruct) aspects of behavior. In particular, you may decode either the AE latents or the ARHMM states on a frame-by-frame basis given the surrounding window of neural activity.

The architecture options consist of a linear model or feedforward neural network: exact architecture parameters such as number of layers in the neural network can be specified in decoding_ae_model.json or decoding_arhmm_model.json. The size of the window of neural activity used to reconstruct each frame of AE latents or ARHMM states is set by n_lags: the neural activity from t-n_lags:t+n_lags will be used to predict the latents or states at time t.

To begin fitting decoding models, copy the example json files decoding_ae_model.json, decoding_arhmm_model.json, decoding_compute.json, and decoding_training.json into your .behavenet directory. cd to the behavenet directory in the terminal, and run:

Decoding ARHMM states:

$: python behavenet/fitting/decoding_grid_search.py --data_config ~/.behavenet/musall_vistrained_params.json --model_config ~/.behavenet/decoding_arhmm_model.json --training_config ~/.behavenet/decoding_training.json --compute_config ~/.behavenet/decoding_compute.json

or

Decoding AE latents:

$: python behavenet/fitting/decoding_grid_search.py --data_config ~/.behavenet/musall_vistrained_params.json --model_config ~/.behavenet/decoding_ae_model.json --training_config ~/.behavenet/decoding_training.json --compute_config ~/.behavenet/decoding_compute.json

It is also possible to decode the motion energy of the AE latents, defined as the absolute value of the difference between neighboring time points; to do so make the following change in the model json: model_class: 'neural-ae-me'

Decoding with subsets of neurons

Continuing with the toy dataset introduced in the data structure documentation, below are some examples for how to modify the decoding data json file to decode from user-specified groups of neurons:

Example 0:

Use all neurons:

{
"subsample_idxs_group_0": null, // not used
"subsample_idxs_group_1": null, // not used
"subsample_idxs_dataset": null, // not used
"subsample_method": "none"      // no subsampling, use all neurons
}

Example 1:

Use the indices in the HDF5 dataset regions/idxs_lr/AUD_L:

{
"subsample_idxs_group_0": "regions", // top-level group
"subsample_idxs_group_1": "idxs_lr", // second-level group
"subsample_idxs_dataset": "AUD_L",   // dataset name
"subsample_method": "single"         // subsample, use single region
}

Example 2:

Fit separate decoders for each dataset of indices in the HDF5 group regions/idxs_lr:

{
"subsample_idxs_group_0": "regions", // top-level group
"subsample_idxs_group_1": "idxs_lr", // second-level group
"subsample_idxs_dataset": "all",     // dataset name
"subsample_method": "single"         // subsample, use single regions
}

In this toy example, these options will fit 4 decoders, each using a different set of indices: AUD_R, AUD_L, VIS_L, and VIS_R.

Note

At this time the option subsample_idxs_dataset can only accept a single string as an argument; therefore you can use all to fit decoders using all datasets in the specified index group, or you can specify a single dataset (e.g. AUD_L in this example). You cannot, for example, provide a list of strings.

Example 3:

Use all indices except those in the HDF5 dataset regions/idxs_lr/AUD_L (“loo” stands for “leave-one-out”):

{
"subsample_idxs_group_0": "regions", // top-level group
"subsample_idxs_group_1": "idxs_lr", // second-level group
"subsample_idxs_dataset": "AUD_L",   // dataset name
"subsample_method": "loo"            // subsample, use all but specified region
}

In this toy example, the combined neurons from AUD_R, VIS_L and VIS_R would be used for decoding (i.e. not the neurons in the specified region AUD_L).

Example 4:

For each dataset in regions/indxs_lr, fit a decoder that uses all indices except those in the dataset:

{
"subsample_idxs_group_0": "regions", // top-level group
"subsample_idxs_group_1": "idxs_lr", // second-level group
"subsample_idxs_dataset": "all",     // dataset name
"subsample_method": "loo"            // subsample, use all but specified region
}

Again referring to the toy example, these options will fit 4 decoders, each using a different set of indices:

  1. AUD_L, VIS_L, and VIS_R (not AUD_R)

  2. AUD_R, VIS_L, and VIS_R (not AUD_L)

  3. AUD_R, AUD_L, and VIS_L (not VIS_R)

  4. AUD_R, AUD_L, and VIS_R (not VIS_L)

Decoding arbitrary covariates

BehaveNet also uses the above decoding infrastructure to allow users to decode an arbitrary set of labels from neural activity; these could be markers from pose estimation software, stimulus information, or other task variables. In order to fit these models, the data HDF5 needs to be augmented to include a new HDF5 group named labels, which contains an HDF5 dataset for each trial. See the data structure documentation for more information.

Once the labels have been added to the data file, you can decode labels as you would CAE latents above; the only changes that are necessary is the addition of the field n_labels in the data json, and changing the model class in the model json from either neural-ae or neural-arhmm to neural-labels.

Note

The current BehaveNet implementation only allows for decoding continuous labels using a Gaussian noise distribution; support for binary and count data forthcoming.

Bayesian decoder

Coming soon!

Advanced user guide

Slurm job submission

Using Behavenet with Slurm is simple: given the slurm submission information in a premade .sh file, test-tube will automatically submit all of the jobs for a given grid search with those slurm settings (see Grid searching with test tube for more details).

Steps

  1. Create an .sh file with the slurm job parameters that you wish to use for this group of jobs (i.e. all the models in the grid search). For example, your .sh file could be:

#!/bin/bash
#SBATCH --job-name=ae_grid_search    # Job name
#SBATCH --mail-type=END,FAIL          # Mail events (NONE, BEGIN, END, FAIL, ALL)
#SBATCH --mail-user=email@gmail.com   # Where to send mail
#SBATCH --time=00:05:00               # Time limit hrs:min:sec
  1. Add slurm hyperparameters (as specified in hyperparameters glossary) to your compute.json

  2. Run the python script as specified throughout these docs and BehaveNet/test-tube will take care of the rest!

Loading a trained model

After you’ve fit one or more models, often you’ll want to load these models and their associated data generator to perform further analyses. BehaveNet provides three methods for doing so:

  • Method 1: load the “best” model from a test-tube experiment

  • Method 2: specify the model version in a test-tube experiment

  • Method 3: specify the model hyperparameters in a test-tube experiment

To illustrate these three methods we’ll use an autoencoder as an example. Let’s assume that we’ve trained 5 convolutional autoencoders with 10 latents, each with a different random seed for initializing the weights, and these have all been saved in the test-tube experiment ae-example.

Method 1: load best model

The first option is to load the best model from ae-example. The “best” model is defined as the one with the smallest loss value computed on validation data. If you set the parameter val_check_interval in the ae training json to a nonzero value before fitting, this information has already been computed and saved in a csv file, so this is a relatively fast option. The following code block shows how to load the best model, as well as the associated data generator, from ae-example.

# imports
from behavenet import get_user_dir
from behavenet.fitting.utils import get_best_model_and_data
from behavenet.fitting.utils import get_expt_dir
from behavenet.fitting.utils import get_lab_example
from behavenet.fitting.utils import get_session_dir
from behavenet.models import AE as Model

# define necessary hyperparameters
hparams = {
    'data_dir': get_user_dir('data'),
    'save_dir': get_user_dir('save'),
    'experiment_name': 'ae-example',
    'model_class': 'ae',
    'model_type': 'conv',
    'n_ae_latents': 10,
}

# programmatically fill out other hparams options
get_lab_example(hparams, 'musall', 'vistrained')
hparams['session_dir'], sess_ids = get_session_dir(hparams)
hparams['expt_dir'] = get_expt_dir(hparams)

# use helper function to load model and data generator
model, data_generator = get_best_model_and_data(hparams, Model, version='best')

Method 2: specify the model version

The next option requires that you know in advance which test-tube version you want to load. In this example, we’ll load version 3. All you need to do is replace version='best' with version=3 in the final line above.

# use helper function to load model and data generator
model, data_generator = get_best_model_and_data(hparams, Model, version=3)

Method 3: specify model hyperparameters

The final option gives you the most control - you can specify all relevant hyperparameters needed to define the model and the data generator, and load that specific model.

# imports
from behavenet import get_user_dir
from behavenet.fitting.utils import experiment_exists
from behavenet.fitting.utils import get_best_model_and_data
from behavenet.fitting.utils import get_expt_dir
from behavenet.fitting.utils import get_lab_example
from behavenet.fitting.utils import get_session_dir
from behavenet.models import AE as Model

# define necessary hyperparameters
hparams = {
    'data_dir': get_user_dir('data'),
    'save_dir': get_user_dir('save'),
    'experiment_name': 'ae-example',
    'model_class': 'ae',
    'model_type': 'conv',
    'n_ae_latents': 10,
    'rng_seed_data': 0,
    'trial_splits': '8;1;1;0',
    'train_frac': 1,
    'rng_seed_model': 0,
    'fit_sess_io_layers': False,
    'learning_rate': 1e-4,
    'l2_reg': 0,
}

# programmatically fill out other hparams options
get_lab_example(hparams, 'musall', 'vistrained')
hparams['session_dir'], sess_ids = get_session_dir(hparams)
hparams['expt_dir'] = get_expt_dir(hparams)

# find the version for these hyperparameters; returns None for version if it doesn't exist
exists, version = experiment_exists(hparams, which_version=True)

# use helper function to load model and data generator
model, data_generator = get_best_model_and_data(hparams, Model, version=version)

You will need to specify the following entries in hparams regardless of the model class:

  • ‘rng_seed_data’

  • ‘trial_splits’

  • ‘train_frac’

  • ‘rng_seed_model’

  • ‘model_class’

  • ‘model_type’

For the autencoder, we need to additionally specify n_ae_latents, fit_sess_io_layers, learning_rate, and l2_reg. Check out the source code for behavenet.fitting.utils.get_model_params() to see which entries are required for other model classes.

Iterating through the data

Below is an example of how to iterate through the data generator and load batches of data:

# select data type to load
dtype = 'train'  # 'train' | 'val' | 'test'

# reset data iterator for this data type
data_generator.reset_iterators(dtype)

# loop through all batches for this data type
for i in range(data_generator.n_tot_batches[dtype]):

    batch, sess = data_generator.next_batch(dtype)
    # "batch" is a dict with keys for the relevant signal, e.g. 'images', 'neural', etc.
    # "sess" is an integer denoting the dataset this batch comes from

    # ... perform analyses ...

Training a model with multiple datasets

The statistical models that comprise BehaveNet - autoencoders, ARHMMs, neural network decoders - often require large amounts of data to avoid overfitting. While the amount of data collected in an hour long experimental session may suffice, every one of these models will benefit from additional data. If data is collected from multiple experimental sessions, and these data are similar enough (e.g. same camera placement/contrast across sessions), then you can train BehaveNet models on all of this data simultaneously.

BehaveNet provides two methods for specifying the experimental sessions used to train a model:

  • Method 1: use all sessions from a specified animal, experiment, or lab

  • Method 2: specify the sessions in a csv file

The first method is simpler, while the second method offers greater control. Both of these methods require modifying the data configuration json before training. We’ll use the Musall dataset as an example; below is the relevant section of the json file located in behavenet/configs/data_default.json that we will modify below.

"lab": "musall", # type: str
"expt": "vistrained", # type: str
"animal": "mSM30", # type: str
"session": "10-Oct-2017", # type: str
"sessions_csv": "", # type: str, help: specify multiple sessions
"all_source": "save", # type: str, help: "save" or "data"

The Musall dataset provided with the repo (see behavenet/example/00_data.ipynb) contains autoencoders trained on two sessions individually, as well as a single autoencoder trained on both sessions as an example of this feature.

Method 1: the “all” keyword

This method is appropriate if you want to fit a model on all sessions from a specified animal, experiment, or lab. For example, if we want to fit a model on all sessions from animal mSM30, we would modify the session parameter value to all:

"lab": "musall", # type: str
"expt": "vistrained", # type: str
"animal": "mSM30", # type: str
"session": "all", # type: str
"sessions_csv": "", # type: str, help: specify multiple sessions
"all_source": "save", # type: str, help: "save" or "data"

In this case the resulting models will be stored in the directory save_dir/musall/vistrained/mSM30/multisession-xx, where xx is selected automatically. BehaveNet will create a csv file named session_info.csv inside the multisession directory that lists the lab, expt, animal, and session for all sessions in that multisession.

If we want to fit a model on all sessions from all animals in the vistrained experiment, we would modify the animal parameter value to all:

"lab": "musall", # type: str
"expt": "vistrained", # type: str
"animal": "all", # type: str
"session": "all", # type: str
"sessions_csv": "", # type: str, help: specify multiple sessions
"all_source": "save", # type: str, help: "save" or "data"

In this case the resulting models will be stored in the directory save_dir/musall/vistrained/multisession-xx. The string value for session does not matter; BehaveNet searches for the all keyword starting at the lab level and moves down; once it finds the all keyword it ignores all further entries.

Note

The all_source parameter in the json file is included to resolve an ambiguity with the “all” keyword. For example, let’s assume you use all at the session level for a single animal. If data for 6 sessions exist for that animal, and BehaveNet models have been fit to 4 of those 6 sessions, then setting "all_source": "data" will use all 6 sessions with data. On the other hand, setting "all_source": "save" will use all 4 sessions that have been previously used to fit models.

Method 2: specify sessions in a csv file

This method is appropriate if you want finer control over which sessions are included; for example, if you want all sessions from one animal, as well as all but one session from another animal. To specify these sessions, you can construct a csv file with the four column headers lab, expt, animal, and session (see below). You can then provide this csv file (let’s say it’s called data_dir/example_sessions.csv) as the value for the sessions_csv parameter:

"lab": "musall", # type: str
"expt": "vistrained", # type: str
"animal": "all", # type: str
"session": "all", # type: str
"sessions_csv": "data_dir/example_sessions.csv", # type: str, help: specify multiple sessions
"all_source": "save", # type: str, help: "save" or "data"

The sessions_csv parameter takes precedence over any values supplied for lab, expt, animal, session, and all_source.

Below is an example csv file that includes two sessions from one animal:

lab,expt,animal,session
musall,vistrained,mSM36,05-Dec-2017
musall,vistrained,mSM36,07-Dec-2017

Here is another example that include the previous two sessions, as well as a third from a different animal:

lab,expt,animal,session
musall,vistrained,mSM30,12-Oct-2017
musall,vistrained,mSM36,05-Dec-2017
musall,vistrained,mSM36,07-Dec-2017

Loading a trained multisession model

The approach is almost identical to that laid out in Loading a trained model; namely, you can either specify the “best” model, the model version, or fully specify all the model hyperparameters. The one necessary change is to alert BehaveNet that you want to load a multisession model. As above, you can do this by either using the “all” keyword or a csv file. The code snippets below illustrate both of these methods when loading the “best” model.

Method 1: use the “all” keyword to specify all sessions for a particular animal:

# imports
from behavenet import get_user_dir
from behavenet.fitting.utils import get_best_model_and_data
from behavenet.fitting.utils import get_expt_dir
from behavenet.fitting.utils import get_lab_example
from behavenet.fitting.utils import get_session_dir
from behavenet.models import AE as Model

# define necessary hyperparameters
hparams = {
    'data_dir': get_user_dir('data'),
    'save_dir': get_user_dir('save'),
    'lab': 'musall',
    'expt': 'vistrained',
    'animal': 'mSM30',
    'session': 'all',  # use all sessions for animal mSM30
    'experiment_name': 'ae-example',
    'model_class': 'ae',
    'model_type': 'conv',
    'n_ae_latents': 10,
}

# programmatically fill out other hparams options
hparams['session_dir'], sess_ids = get_session_dir(hparams)
hparams['expt_dir'] = get_expt_dir(hparams)

# use helper function to load model and data generator
model, data_generator = get_best_model_and_data(hparams, Model, version='best')

As above, the all keyword can also be used at the animal or expt level, though not currently at the lab level.

Method 2: use a sessions csv file:

# imports
from behavenet import get_user_dir
from behavenet.fitting.utils import get_best_model_and_data
from behavenet.fitting.utils import get_expt_dir
from behavenet.fitting.utils import get_lab_example
from behavenet.fitting.utils import get_session_dir
from behavenet.models import AE as Model

# define necessary hyperparameters
hparams = {
    'data_dir': get_user_dir('data'),
    'save_dir': get_user_dir('save'),
    'sessions_csv': '/path/to/csv/file',
    'experiment_name': 'ae-example',
    'model_class': 'ae',
    'model_type': 'conv',
    'n_ae_latents': 10,
}

# programmatically fill out other hparams options
hparams['session_dir'], sess_ids = get_session_dir(hparams)
hparams['expt_dir'] = get_expt_dir(hparams)

# use helper function to load model and data generator
model, data_generator = get_best_model_and_data(hparams, Model, version='best')

In both cases, iterating through the data proceeds exactly as when using a single session, and the second return value from data_generator.next_batch() identifies which session the batch belongs to.

PS-VAE hyperparameter search guide

The PS-VAE objective function \(\mathscr{L}_{\text{PS-VAE}}\) is comprised of several different terms:

\[\mathscr{L}_{\text{PS-VAE}} = \mathscr{L}_{\text{frames}} + \alpha \mathscr{L}_{\text{labels}} + \mathscr{L}_{\text{KL-s}} + \mathscr{L}_{\text{ICMI}} + \beta \mathscr{L}_{\text{TC}} + \mathscr{L}_{\text{DWKL}}\]

where

  • \(\mathscr{L}_{\text{frames}}\): log-likelihood of the video frames

  • \(\mathscr{L}_{\text{labels}}\): log-likelihood of the labels

  • \(\mathscr{L}_{\text{KL-s}}\): KL divergence of the supervised latents

  • \(\mathscr{L}_{\text{ICMI}}\): index-code mutual information of the unsupervised latents

  • \(\mathscr{L}_{\text{TC}}\): total correlation of the unsupervised latents

  • \(\mathscr{L}_{\text{DWKL}}\): dimension-wise KL of the unsupervised latents

There are two important hyperparameters of the model that we address below: \(\alpha\), which weights the reconstruction of the labels, and \(\beta\), which weights the factorization of the unsupervised latent space. The purpose of this guide is to propose a series of model fits that efficiently explores this space of hyperparameters, as well as point out several BehaveNet plotting utilities to assist in this exploration.

How to select \(\alpha\)

The hyperparameter \(\alpha\) controls the strength of the label log-likelihood term, which needs to be balanced against the frame log-likelihood term. We first recommend z-scoring each individual label, which removes the scale of the labels as a confound. We then recommend fitting models with a range of \(\alpha\) values, while setting the defaults \(\beta=1\) (no extra weight on the total correlation term). In our experience the range \(\alpha=[50, 100, 500, 1000]\) is a reasonable range to start with. The “best” value for \(\alpha\) is subjective because it involves a tradeoff between pixel log-likelihood (or the related mean square error, MSE) and label log-likelihood (or MSE). After choosing a suitable value, we will fix \(\alpha\) and vary \(\beta\).

How to select \(\beta\)

The choice of \(\beta\) is more difficult because there does not yet exist a single robust measure of “disentanglement” that can tell us which models learn a suitable unsupervised representation. Instead we will fit models with a range of hypeparameters, then use a quantitative metric to guide a qualitative analysis.

A reasonable range to start with is \(\beta=[1, 5, 10, 20]\). How, then, do we choose the “best” value of \(\beta\)? Currently our best advice is to compute the correlation of the training data across all pairs of unsupervised dimensions. The value of \(\beta\) that minimizes the average of the pairwise correlations is a good place to start more qualitative evaluations.

Ultimately, the choice of the “best” model comes down to a qualitative evaluation, the latent traversal. A latent traversal is the result of changing the value of a latent dimension while keeping the value of all other latent dimensions fixed. If the model has learned an interpretable representation then the resulting generated frames should show one single behavioral feature changing per dimension - an arm, or a jaw, or the chest (see below for more information on tools for constructing and visualizing these traversals). In order to choose the “best” model, we perform these latent traversals for all values of \(\beta\) and look at the resulting latent traversal outputs. The model with the (subjectively) most interpretable dimensions is then chosen.

A note on model robustness

We have found the PS-VAE to be somewhat sensitive to initialization of the neural network parameters. We also recommend choosing the set of hyperparamters with the lowest pairwise correlations and refitting the model with several random seeds (by changing the rng_seed_model parameter of the ae_model.json file), which may lead to even better results.

Tools for investigating PS-VAE model fits

The functions listed below are provided in the BehaveNet plotting module (behavenet.plotting) to facilitate model checking and comparison at different stages.

Hyperparameter search visualization

The function behavenet.plotting.cond_ae_utils.plot_hyperparameter_search_results() creates a variety of diagnostic plots after the user has performed the \(\alpha\) search and the \(\beta\) search detailed above:

  • pixel mse as a function of \(\alpha\), num latents (for fixed \(\beta\))

  • label mse as a function of \(\alpha\), num_latents (for fixed \(\beta\))

  • pixel mse as a function of \(\beta\) (for fixed \(\alpha\), n_ae_latents)

  • label mse as a function of \(\beta\) (for fixed \(\alpha\), n_ae_latents)

  • index-code mutual information (part of the KL decomposition) as a function of \(\beta\) (for fixed \(\alpha\), n_ae_latents)

  • total correlation (part of the KL decomposition) as a function of \(\beta\) (for fixed \(\alpha\), n_ae_latents)

  • dimension-wise KL (part of the KL decomposition) as a function of \(\beta\) (for fixed \(\alpha\), n_ae_latents)

  • average correlation coefficient across all pairs of unsupervised latent dims as a function of \(\beta\) (for fixed \(\alpha\), n_ae_latents)

These plots help with the selection of hyperparameter settings.

Model training curves

The function behavenet.plotting.cond_ae_utils.plot_psvae_training_curves() creates training plots for each term in the PS-VAE objective function for a single model:

  • total loss

  • pixel mse

  • label \(R^2\) (note the objective function contains the label MSE, but \(R^2\) is easier to parse)

  • KL divergence of supervised latents

  • index-code mutual information of unsupervised latents

  • total correlation of unsupervised latents

  • dimension-wise KL of unsupervised latents

A function argument allows the user to plot either training or validation curves. These plots allow the user to check whether or not models have trained to completion.

Label reconstruction

The function behavenet.plotting.cond_ae_utils.plot_label_reconstructions() creates a series of plots that show the true labels and their PS-VAE reconstructions for a given list of batches. These plots are useful for qualitatively evaluating the supervised subspace of the PS-VAE; a quantitative evaluation (the label MSE) can be found in the metrics.csv file created in the model folder during training.

Latent traversals: plots

The function behavenet.plotting.cond_ae_utils.plot_latent_traversals() displays video frames representing the traversal of chosen dimensions in the latent space. This function uses a single base frame to create all traversals.

Latent traversals: movies

The function behavenet.plotting.cond_ae_utils.make_latent_traversal_movie() creates a multi-panel movie with each panel showing traversals of an individual latent dimension. The traversals will start at a lower bound, increase to an upper bound, then return to a lower bound; the traversal of each dimension occurs simultaneously. It is also possible to specify multiple base frames for the traversals; the traversal of each base frame is separated by several blank frames.

MSPS-VAE hyperparameter search guide

The MSPS-VAE objective function \(\mathscr{L}_{\text{MSPS-VAE}}\) is comprised of several different terms:

\[\mathscr{L}_{\text{PS-VAE}} = \mathscr{L}_{\text{frames}} + \alpha \mathscr{L}_{\text{labels}} + \mathscr{L}_{\text{KL-s}} + \mathscr{L}_{\text{ICMI}} + \beta \mathscr{L}_{\text{TC}} + \mathscr{L}_{\text{DWKL}} + \delta \mathscr{L}_{\text{triplet}}\]

where

  • \(\mathscr{L}_{\text{frames}}\): log-likelihood of the video frames

  • \(\mathscr{L}_{\text{labels}}\): log-likelihood of the labels

  • \(\mathscr{L}_{\text{KL-s}}\): KL divergence of the supervised latents

  • \(\mathscr{L}_{\text{ICMI}}\): index-code mutual information of the unsupervised latents

  • \(\mathscr{L}_{\text{TC}}\): total correlation of the unsupervised latents

  • \(\mathscr{L}_{\text{DWKL}}\): dimension-wise KL of the unsupervised latents

  • \(\mathscr{L}_{\text{triplet}}\): triplet loss on background latents

There are three important hyperparameters of the model that we address below: \(\alpha\), which weights the reconstruction of the labels; \(\beta\), which weights the factorization of the unsupervised latent space; and \(\delta\), which weights the triplet loss of the background latent space (this pushes data points from the same session togther, and data points from different sessions farther apart). The purpose of this guide is to propose a series of model fits that efficiently explores this space of hyperparameters, as well as point out several BehaveNet plotting utilities to assist in this exploration.

How to select \(\alpha\)

The hyperparameter \(\alpha\) controls the strength of the label log-likelihood term, which needs to be balanced against the frame log-likelihood term. We first recommend z-scoring each individual label within each session, which removes the scale of the labels as a confound. We then recommend fitting models with a range of \(\alpha\) values, while setting the defaults \(\beta=1\) (no extra weight on the total correlation term) and \(\delta=50\) (slight penalty on triplet loss). In our experience the range \(\alpha=[50, 100, 500, 1000]\) is a reasonable range to start with. The “best” value for \(\alpha\) is subjective because it involves a tradeoff between pixel log-likelihood (or the related mean square error, MSE) and label log-likelihood (or MSE). After choosing a suitable value, we will fix \(\alpha\) and vary \(\beta\) and \(\delta\).

How to select \(\beta\) and \(\delta\)

The choice of \(\beta\) and \(\delta\) is more difficult because there does not yet exist a single robust measure of “disentanglement” that can tell us which models learn a suitable unsupervised representation. Instead we will fit models with a range of hypeparameters, then use a quantitative metric to guide a qualitative analysis.

A reasonable range to start with is \(\beta=[1, 5, 10, 20]\) and \(\delta=[50, 100, 500, 1000]\). How, then, do we choose good values for \(\beta\) and \(\delta\)? Currently our best advice is to compute the correlation of the training data across all pairs of unsupervised dimensions. The values of \(\beta\) and \(\delta\) that minimizes the average of the pairwise correlations is a good place to start more qualitative evaluations. We also use another metric here to assist in the selection process: the accuracy of a classifier trained to predict session identity from the unsupervised representation. This metric should be near chance (1 / number of sessions) for models that have successfully learned to place inter-session variability in the background subspace. This metric may be heavily influenced by the value of \(\delta\), and thus can provide a clearer distinction between models.

Ultimately, the choice of the “best” model comes down to a qualitative evaluation, the latent traversal. A latent traversal is the result of changing the value of a latent dimension while keeping the value of all other latent dimensions fixed. If the model has learned an interpretable representation then the resulting generated frames should show one single behavioral feature changing per dimension - an arm, or a jaw, or the chest (see below for more information on tools for constructing and visualizing these traversals). In order to choose the “best” model, we perform these latent traversals for all values of \(\beta\) and \(\delta\) and look at the resulting latent traversal outputs. The model with the (subjectively) most interpretable dimensions is then chosen.

A note on model robustness

We have found the MSPS-VAE to be somewhat sensitive to initialization of the neural network parameters. We also recommend choosing the set of hyperparamters with the lowest pairwise correlations and refitting the model with several random seeds (by changing the rng_seed_model parameter of the ae_model.json file), which may lead to even better results.

Tools for investigating MSPS-VAE model fits

The functions listed below are provided in the BehaveNet plotting module (behavenet.plotting) to facilitate model checking and comparison at different stages.

Hyperparameter search visualization

The function behavenet.plotting.cond_ae_utils.plot_mspsvae_hyperparameter_search_results() creates a variety of diagnostic plots after the user has performed the \(\alpha\) search and the \(\beta/\delta\) search detailed above:

  • pixel mse as a function of \(\alpha\), num latents (for fixed \(\beta, \delta\))

  • label mse as a function of \(\alpha\), num_latents (for fixed \(\beta, \delta\))

  • pixel mse as a function of \(\beta, \delta\) (for fixed \(\alpha\), n_ae_latents)

  • label mse as a function of \(\beta, \delta\) (for fixed \(\alpha\), n_ae_latents)

  • index-code mutual information (part of the KL decomposition) as a function of \(\beta, \delta\) (for fixed \(\alpha\), n_ae_latents)

  • total correlation (part of the KL decomposition) as a function of \(\beta, \delta\) (for fixed \(\alpha\), n_ae_latents)

  • dimension-wise KL (part of the KL decomposition) as a function of \(\beta, \delta\) (for fixed \(\alpha\), n_ae_latents)

  • average correlation coefficient across all pairs of unsupervised latent dims as a function of \(\beta, \delta\) (for fixed \(\alpha\), n_ae_latents)

  • session classification accuracy as a function of \(\beta, \delta\) (for fixed \(\alpha\), n_ae_latents)

  • triplet loss as a function of \(\beta, \delta\) (for fixed \(\alpha\), n_ae_latents)

These plots help with the selection of hyperparameter settings.

Model training curves

The function behavenet.plotting.cond_ae_utils.plot_mspsvae_training_curves() creates training plots for each term in the PS-VAE objective function for a single model:

  • total loss

  • pixel mse

  • label R^2 (note the objective function contains the label MSE, but R^2 is easier to parse)

  • KL divergence of supervised latents

  • index-code mutual information of unsupervised latents

  • total correlation of unsupervised latents

  • dimension-wise KL of unsupervised latents

  • triplet loss

A function argument allows the user to plot either training or validation curves. These plots allow the user to check whether or not models have trained to completion.

Label reconstruction

The function behavenet.plotting.cond_ae_utils.plot_label_reconstructions() creates a series of plots that show the true labels and their MSPS-VAE reconstructions for a given list of batches. These plots are useful for qualitatively evaluating the supervised subspace of the MSPS-VAE; a quantitative evaluation (the label MSE) can be found in the metrics.csv file created in the model folder during training.

Latent traversals: plots

The function behavenet.plotting.cond_ae_utils.plot_latent_traversals() displays video frames representing the traversal of chosen dimensions in the latent space. This function uses a single base frame to create all traversals.

Latent traversals: movies

The function behavenet.plotting.cond_ae_utils.make_latent_traversal_movie() creates a multi-panel movie with each panel showing traversals of an individual latent dimension. The traversals will start at a lower bound, increase to an upper bound, then return to a lower bound; the traversal of each dimension occurs simultaneously. It is also possible to specify multiple base frames for the traversals; the traversal of each base frame is separated by several blank frames.

BehaveNet data structure

Introduction

In order to quickly and easily fit many models, BehaveNet uses a standardized data structure. “Raw” experimental data such as behavioral videos and (processed) neural data are stored in the HDF5 file format. This file format can accomodate large and complex datasets, and is easy to work with thanks to a high-level python API.

HDF is an acronym for Hierarchical Data Format, and one can think of it like a full directory structure inside of a single file. HDF5 “groups” are analogous to directories, while HDF5 “datasets” are analogous to files. The BehaveNet code uses up to 3 HDF5 groups: images, masks (for masking images), and neural. Each of these HDF5 groups contains multiple HDF5 datasets - one for each experimental trial. These datasets should have names that follow the pattern trial_%04i - datasets with more than 10000 trials are not currently supported with this naming convention.

BehaveNet models are trained on batches of data, which here are defined as one trial per batch. For datasets that do not have a trial structure (i.e. spontaneous behavior) we recommend splitting frames into arbitrarily defined “trials”, the length of which should depend on the autocorrelation of the behavior (i.e. trials should not be shorter than the temporal extent of relevant behaviors). For the NP dataset in the original paper we used batch sizes of 1000 frames (~25 sec), and inserted additional gap trials between training, validation, and testing trials to minimize the possibility that good model performance was due to similarity of trials.

Below is a sample python script demonstrating how to create an HDF5 file with video data and neural data. Video data is assumed to be in a list, where each list element corresponds to a single trial, and is a numpy array of shape (n_frames, n_channels, y_pix, x_pix). Neural data is assumed to be in the same format; a corresponding list of numpy arrays of shape (n_frames, n_neurons). BehaveNet does not require all trials to be of the same length, but does require that for each trial the images and neural activity have the same number of frames. This may require you to interpolate/bin video or neural data differently than the rate at which it was acquired.

Notes:

  • for large experiments, having all of this (video) data in memory to create the HDF5 file might be infeasible, and more sophisticated processing will be required

  • neural data is only required for fitting decoding models; it is still possible to fit autoencoders and ARHMMs when the HDF5 file only contains images

  • masks should be the same size as images; a value of 0 excludes the pixel from the loss function, a value of 1 includes it

  • the python package h5py is required for creating the HDF5 file, and is automatically installed with the BehaveNet package.

import h5py

# assume images are in an np array named "images_np"; this should be a an array of dtype
# 'uint8', and values should be between 0 and 255. The data type is converted to float
# and values are divided by 255 during the data loading process.

# assume neural activity is in an np array named "neural_np"; this can be spike count data
# or continuous-valued data for e.g. calcium imaging experiments

hdf5_file = '/path/to/data.hdf5'  # path needs to exist, but not 'data.hdf5' file

with h5py.File(hdf5_file, 'w', libver='latest', swmr=True) as f:

    # enable single write, multi-read - needed for simultaneous model fitting
    f.swmr_mode = True

    # create "image" HDF5 group
    group_i = f.create_group('images')

    # create "neural" HDF5 group
    group_n = f.create_group('neural')

    # create a dataset for each trial within groups
    for trial in range(len(images_np)):

        # create dataset in "image" group
        # images_np[trial] should be of shape (n_frames, n_channels, y_pix, x_pix)
        group_i.create_dataset('trial_%04i' % trial, data=images_np[trial], dtype='uint8')

        # create dataset in "neural" group
        # neural_np[trial] should be of shape (n_frames, n_neurons)
        group_n.create_dataset('trial_%04i' % trial, data=neural_np[trial], dtype='float32')

A more in-depth example can be found in the function behavenet.data.preprocess.build_hdf5().

Identifying subsets of neurons

It is possible that the neural data used for encoding and decoding models will have natural partitions - for example, neurons belonging to different brain regions or cell types. In this case you may be interested in, say, decoding behavior from each brain region individually, as well as all together. BehaveNet provides this capability through the addition of another HDF5 group. This group can have any name, but for illustration purposes we will use the name “regions” (this name will be later be provided in the updated data json file).

The “regions” group contains a second level of (again user-defined) groups, which will define different index groupings. As a concrete example, let’s say we have neural data with 100 neurons:

  • indices 00-24 are neurons in left auditory cortex

  • indices 25-49 are neurons in right auditory cortex

  • indices 50-74 are neurons in left visual cortex

  • indices 75-99 are neurons in right visual cortex

We will define this “grouping” of indices in a python dict:

neural_idxs_lr = {
    'AUD_L': np.arange(0, 25),
    'AUD_R': np.arange(25, 50),
    'VIS_L': np.arange(50, 75),
    'VIS_R': np.arange(75, 100)
}

We can also define another “grouping” of indices that ignores hemisphere information:

neural_idxs = {
    'AUD': np.arange(0, 50),
    'VIS': np.arange(50, 100)
}

We can then store these indices in the data HDF5 by modifying the above script:

...

# create "neural" HDF5 group
group_n = f.create_group('neural')

# create "regions" HDF5 group
group_r0 = f.create_group('regions')

# create "idxs_lr" HDF5 group inside the "regions" group
group_r1a = group_r0.create_group('idxs_lr')
# insert the index info into datasets inside the regions/idxs_lr group
for region_name, region_idxs in neural_idxs_lr.items():
    group_r1a.create_dataset(region_name, data=region_idxs)

# create "idxs" HDF5 group inside the "regions" group
group_r1b = group_r0.create_group('idxs')
# insert the index info into datasets inside the regions/idxs group
for region_name, region_idxs in neural_idxs.items():
    group_r1b.create_dataset(region_name, data=region_idxs)

# create a dataset for each trial within groups
for trial in range(len(images_np)):

...

This HDF5 file will now have the following addtional datasets:

  • regions/idxs_lr/AUD_L

  • regions/idxs_lr/AUD_R

  • regions/idxs_lr/VIS_L

  • regions/idxs_lr/VIS_R

  • regions/idxs/AUD

  • regions/idxs/VIS

Just as the top-level group (here named “regions”) can have an arbitrary name (later specified in the data json file), the second-level groups (here named “idxs_lr” and “idxs”) can also have arbitrary names, and there can be any number of them, as long as the datasets within them contain valid indices into the neural data. The specific set of indices used for any analyses will be specified in the data json file. See the decoding documentation for an example of how to decode behavior using specified subsets of neurons.

Including labels for ARHMMs and conditional autoencoders

In order to fit conditional autoencoder models, you will need to include additional information about labels in the HDF5 file. These labels can be outputs from pose estimation software, or other behavior-related signals such as pupil diameter or lick times. These labels should be stored in an HDF5 group named labels. As before, the labels group contains multiple HDF5 datasets - one for each experimental trial. These datasets should also follow the pattern trial_%04i, and match the image data in the corresponding image dataset images/trial_%04i. If the image data in a given trial is of shape (n_frames, n_channels, y_pix, x_pix), then the corresponding label data should be of shape (n_frames, n_markers). Note that, when using pose estimation software, each marker has an x- and y-coordinate, so tracking four body parts will result in an 8-dimensional set of labels.

It is also possible to fit ARHMMs directly to labels rather than the outputs of an autoencoder. In this case labels is the only necessary HDF5 group, though including a corresponding images group will allow you to utilize more of the ARHMM visualization tools. To fit an ARHMM on label data, you simply need to change the model_class entry of the arhmm model json from arhmm to arhmm-labels (see the json config arhmm_labels_model.json).

Note

The matrix subspace projection model implemented in BehaveNet learns a linear mapping from the original latent space to the predicted labels that does not contain a bias term. Therefore you should center each label before adding them to the HDF5 file. Additionally, normalizing each label by its standard deviation can make searching across msp weights less dependent on the size of the input image.

Hyperparameter glossary

The BehaveNet code requires a diverse array of hyperparameters (hparams) to specify details about the data, computational resources, training algorithms, and the models themselves. This glossary contains a brief description for each of the hparams options. See the example json files for reasonable hparams defaults.

Data

  • data_dir (str): absolute path to data directory

  • save_dir (str): absolute path to save directory, where models are stored

  • lab (str): lab id

  • expt (str): experiment id

  • animal (str): animal id

  • session (str): session id

  • sessions_csv (str): path to csv file that contains a list of sessions to use for model fitting. The 4 column headers in the csv should be lab, expt, animal, session. If this is not an empty string, it supercedes the information provided in the lab, expt, animal, and session fields above.

  • all_source (str): one of the expt, animal, or session fields can optionally be set to the string "all". For example, if expt is "all", then for the specified lab all sessions from all experiments/animals are collected and fit with the same model. If animal is "all", then for the specified lab and expt all sessions are collected. The field all_source tells the code where to look for the corresponding sessions: "data" will search for all sessions in data_dir; "save" will search for all sessions in save_dir, and as such will only find the sessions that have been previously used to fit models.

  • n_input_channels (str): number of colors channel/camera views in behavioral video

  • y_pixels (int): number of behavioral video pixels in y dimension

  • x_pixels (int): number of behavioral video pixels in x dimension

  • use_output_mask (bool): True` to apply frame-wise output masks (must be a key masks in data HDF5 file)

  • use_label_mask (bool): True` to apply frame-wise masks to labels in conditional ae models (must be a key labels_masks in data HDF5 file)

  • n_labels (bool): specify number of labels when model_class is ‘neural-labels’ or ‘labels-neural’

  • neural_bin_size (float): bin size of neural/video data (ms)

  • neural_type (str): ‘spikes’ | ‘ca’

  • approx_batch_size (str): approximate batch size (number of frames) for gpu memory calculation

For encoders/decoders, additional information can be supplied to control which subsets of neurons are used for encoding/decoding. See the data structure documentation for detailed instructions on how to incorporate this information into your HDF5 data file. The following options must be added to the data json file (an example can be found here):

  • subsample_idxs_group_0 (str): name of the top-level HDF5 group that contains index groups

  • subsample_idxs_group_1 (str): name of the second-level HDF5 group that contains index datasets

  • subsample_idxs_dataset (str): use “all” to have test tube loop over each index dataset in subsample_idxs_group_0/subsample_idxs_group_1, or specify a single user-defined index dataset as a string

  • subsample_method (str): determines how different index datasets are subsampled

    • ‘none’: no subsampling; all neural data is used for encoding/decoding

    • ‘single’: for the index dataset specified by ‘subsample_idxs_dataset’, use just these indices for decoding

    • ‘loo’: leave-one-out; for the index dataset specified by ‘subsample_idxs_dataset’, use all except this dataset for decoding

Computational resources

  • device (str): where to fit pytorch models; ‘cpu’ | ‘cuda’

  • n_parallel_gpus (int): number of gpus to use per model, currently only implemented for AEs

  • tt_n_gpu_trials (int): total number of hyperparameter combinations to fit with test-tube on gpus

  • tt_n_cpu_trials (int): total number of hyperparameter combinations to fit with test-tube on cpus

  • tt_n_cpu_workers (int): total number of cpu cores to use with test-tube for hyperparameter searching

  • mem_limit_gb (float): maximum size of gpu memory; used to filter out randomly generated CAEs that are too large

If using machine without slurm:

  • gpus_viz (str): which gpus are visible to test-tube; multiple gpus are identified as e.g. ‘0;1;4’

If using slurm:

  • slurm (bool): true if using slurm, false otherwise

  • slurm_log_path (str): directory in which to save slurm outputs (.err/.out files)

  • slurm_param_file (str): file name of the .sh file with slurm params for job submission

Training

All models:

  • as_numpy (bool): True to load data as numpy arrays, False to load as pytorch tensors

  • batch_load (bool): True to load data one batch at a time, False to load all data into memory (the data is still served to models in batches)

  • rng_seed_data (int): control randomness when splitting data into train, val, and test trials

  • train_frac (float): if 0 < train_frac < 1.0, defines the fraction of assigned training trials to actually use; if train_frac > 1.0, defines the number of assigned training trials to actually use (rounded to the nearest integer)

  • trial_splits (str): determines number of train/val/test/gap trials; entered as 8;1;1;0, for example. See behavenet.data.data_generator.split_trials() for how these values are used.

  • export_train_plots (bool): True to automatically export training/validation loss as a function of epoch upon completion of training [AEs and ARHMMs only]

  • export_latents (bool): True to automatically export train/val/test autoencoder latents using best model upon completion of training [analogous parameters export_states and export_predictions exist for arhmms and decoders, respectively)

  • rng_seed_train (int): control randomness in batching data

Pytorch models (all but ‘arhmm’ and ‘bayesian-decoding’):

  • val_check_interval: (float): frequency with which metrics are calculated on validation data. These metrics are logged in a csv file via test-tube, and can also be used for early stopping if enabled. If 0 < val_check_interval < 1.0, metrics are computed multiple times per epoch (val_check_interval=0.5 corresponds to checking every half epoch); if val_check_interval > 1.0, defines number of epochs between metric computation.

  • learning_rate (float): learning rate of adam optimizer

  • max_n_epochs (int): maximum number of training epochs

  • min_n_epochs (int): minimum number of training epochs, even when early stopping is used

  • enable_early_stop (bool): if False, training proceeds until maximum number of epochs is reached

  • early_stop_history (int): number of epochs over which to average validation loss

ARHMM:

  • n_iters (int): number of EM iterations (currently no early stopping)

  • arhmm_es_tol (float): relative tolerance for early stopping; training terminates if the absolute value of the difference between the previous log likelihood (ll) and current ll, divided by the current ll, is less than this value

Models

All models:

  • experiment_name (str): name of the test-tube experiment

  • rng_seed_model (int): control initialization of model parameters

  • model_class: (str): name of the model class

    • ‘ae’: autoencoder

    • ‘vae’: variational autoencoder

    • ‘beta-tcvae’: variational autoencoder with beta tc-vae decomposition of elbo

    • ‘cond-ae’: conditional autoencoder

    • ‘cond-vae’: conditional variational autoencoder

    • ‘cond-ae-msp’: autoencoder with matrix subspace projection loss

    • ‘ps-vae’: partitioned subspace variational autoencoder

    • ‘msps-vae’: multi-session partitioned subspace variational autoencoder

    • ‘hmm’: hidden Markov model

    • ‘arhmm’: autoregressive hidden Markov model

    • ‘neural-ae’: decode AE latents from neural activity

    • ‘neural-ae-me’: decode motion energy of AE latents (absolute value of temporal difference) from neural activity

    • ‘neural-arhmm’: decode arhmm states from neural activity

    • ‘ae-neural’: predict neural activity from AE latents

    • ‘arhmm-neural’: predict neural activity from arhmm states

    • ‘labels-images’: decode images from labels with a convolutional decoder

    • ‘bayesian-decoding’: baysian decoding of AE latents and arhmm states from neural activity

Pytorch models (all but ‘arhmm’ and ‘bayesian-decoding’):

  • l2_reg (float): L2 regularization value applied to all model weights

Autoencoder

  • model_type (str): ‘conv’ | ‘linear’

  • n_ae_latents (int): output dimensions of AE encoder network

  • fit_sess_io_layers (bool): True to fit session-specific input and output layers; all other layers are shared across all sessions

  • ae_arch_json (str): null to use the default convolutional autoencoder architecture from the original behavenet paper; otherwise, a string that defines the path to a json file that defines the architecture. An example can be found here.

Variational autoencoders

In addition to the autoencoder parameters defined above,

  • vae.beta (float): weight on KL divergence term in VAE ELBO

  • vae.beta_anneal_epochs (int): number of epochs over which to linearly increase VAE beta

  • beta_tcvae.beta (float) weight on total correlation term in Beta TC-VAE ELBO

  • beta_tcvae.beta_anneal_epochs (int): number of epochs over which to linearly increase Beta TC-VAE beta

  • ps_vae.alpha (float) weight on label reconstruction term in (MS)PS-VAE ELBO

  • ps_vae.beta (float) weight on unsupervised disentangling term in (MS)PS-VAE ELBO

  • ps_vae.delta (float): weight on triplet loss for the background subspace of the MSPS-VAE

  • ps_vae.anneal_epochs (int): number of epochs over which to linearly increase (MS)PS-VAE beta

  • n_background (int): dimensionality of background subspace in the MSPS-VAE

  • n_sessions_per_batch (int): number of sessions to use for a single batch when training the MSPS-VAE (triplet loss needs data from multiple sessions)

Conditional autoencoders

In addition to the autoencoder parameters defined above,

  • conditional_encoder (bool): True to condition encoder on labels when model class is ‘cond-ae’

  • msp.alpha (float): weight on label reconstruction term when model class is ‘cond-ae-msp’

ARHMM

  • model_type (NoneType): not used for ARHMMs

  • n_arhmm_lags (int): number of autoregressive lags (order of AR process)

  • noise_type (str): observation noise; ‘gaussian’ | ‘studentst’ | ‘diagonal_gaussian’ | ‘diagonal_studentst’

  • transitions (float): transition model; ‘stationary’ | ‘sticky’ | ‘recurrent’ | ‘recurrent_only’

  • kappa (float): stickiness parameter that biases diagonal of Markov transition matrix, which increases average state durations

  • ae_experiment_name (str): name of AE test-tube experiment

  • ae_version (str or int): ‘best’ to choose best version in AE experiment, otherwise an integer specifying test-tube version number

  • ae_model_class (str): ‘ae’ | ‘vae’ | ‘beta-tcvae’ | …

  • ae_model_type (str): ‘conv’ | ‘linear’

  • n_ae_latents (int): number of autoencoder latents; this will be the observation dimension in the ARHMM

  • export_train_plots (’bool): True to automatically export training/validation log probability as a function of epoch upon completion of training

  • export_states (bool): True to automatically export train/val/test states using best model upon completion of training

Decoder

For both continuous and discrete decoders:

  • model_type:

    • ‘mlp’ - standard feedforward neural network; use n_hid_layers=0 (see below) for linear regression

    • ‘mlp-mv’ - use the neural network to estimate both the mean and the covariance matrix of the AE latents

    • ‘lstm’ - currently not implemented

  • n_hid_layers (int): number of hidden layers in decoder, not counting data or output layer

  • n_hid_units (int): number of units in all hidden layers; the code will automatically choose the correct number of units for the output layer based on the data size

  • n_lags (int): number of time lags in neural activity to use in predicting outputs; if n_lags=n, then the window of neural activity t-n:t+n is used to predict the outputs at time t (and therefore 2n+1 total time points are used to predict each time point)

  • n_max_lags (int): maximum number of lags the user thinks they may search over; the first n_max_lags and final n_max_lags time points of each batch are not used in the calculation of metrics to make models with differing numbers of lags directly comparable

  • activation (str): activation function of hidden layers; activation function of final layer is automatically chosen based on decoder/data type; ‘linear’ | ‘relu’ | ‘lrelu’ | ‘sigmoid’ | ‘tanh’

  • export_predictions (bool): True to automatically export train/val/test predictions using best model upon completion of training

For the continuous decoder:

  • ae_experiment_name (str): name of AE test-tube experiment

  • ae_version (str or int): ‘best’ to choose best version in AE experiment, otherwise an integer specifying test-tube version number

  • ae_model_class (str): ‘ae’ | ‘vae’ | ‘beta-tcvae’ | …

  • ae_model_type (str): ‘conv’ | ‘linear’

  • n_ae_latents (int): number of autoencoder latents; this will be the dimension of the data predicted by the decoder

  • ae_multisession (int): use if loading latents from an AE that was trained on multiple datasets

For the discrete decoder:

  • n_ae_latents (int): number of autoencoder latents that the ARHMM was trained on

  • ae_model_class (str): ‘ae’ | ‘vae’ | ‘beta-tcvae’ | …

  • ae_model_type (str): ‘conv’ | ‘linear’

  • arhmm_experiment_name (str): name of ARHMM test-tube experiment

  • n_arhmm_states (int): number of ARHMM discrete states; this will be the number of classes the decoder is trained on

  • n_arhmm_lags (int): number of autoregressive lags (order of AR process)

  • kappa (float): ‘kappa’ parameter of the desired ARHMM

  • noise_type (str): ‘noise_type’ parameter of the desired ARHMM; ‘gaussian’ | ‘studentst’

  • arhmm_version (str or int): ‘best’ to choose best version in ARHMM experiment, otherwise an integer specifying test-tube version number

  • arhmm_multisession (int): use if loading states from an ARHMM that was trained on multiple datasets

Bayesian decoder

TODO

BehaveNet API

behavenet.data

Data handing documentation.

behavenet.data.data_generator Module

Classes for splitting and serving data to models.

The data generator classes contained in this module inherit from the torch.utils.data.Dataset class. The user-facing class is the ConcatSessionsGenerator, which can manage one or more datasets. Each dataset is composed of trials, which are split into training, validation, and testing trials using the split_trials(). The default data generator can handle the following data types:

  • images: individual frames of the behavioral video

  • masks: binary mask for each frame

  • labels: i.e. DLC labels

  • neural activity

  • AE latents

  • AE predictions: predictions of AE latents from neural activity

  • ARHMM states

  • ARHMM predictions: predictions of ARHMM states from neural activity

Please see the online documentation at Read the Docs for detailed examples of how to use the data generators.

Functions

split_trials(n_trials[, rng_seed, train_tr, …])

Split trials into train/val/test blocks.

Classes

SingleSessionDatasetBatchedLoad(data_dir[, …])

Dataset class for a single session with batch loading of data.

SingleSessionDataset(data_dir[, lab, expt, …])

Dataset class for a single session.

ConcatSessionsGenerator(data_dir, ids_list)

Dataset class for multiple sessions.

ConcatSessionsGeneratorMulti(data_dir, ids_list)

Dataset class for multiple sessions, which returns multiple sessions per training batch.

behavenet.data.preprocess Module

Utility functions for automatically constructing hdf5 files.

Functions

build_hdf5(save_file, video_file[, …])

Build Behavenet-style HDF5 file from video file and optional label file.

load_raw_labels(file_path, pose_algo[, …])

Load labels and build masks from a variety of standardized source files.

resize_labels(labels, xpix_new, ypix_new, …)

Update label values to reflect scale of corresponding images.

get_frames_from_idxs(cap, idxs)

Helper function to load video segments.

behavenet.data.transforms Module

Tranform classes to process data.

Data generator objects can apply these transforms to batches of data upon loading.

Classes

BlockShuffle(rng_seed)

Shuffle blocks of contiguous discrete states within each trial.

ClipNormalize(clip_val)

Clip upper level of signal and divide by clip value.

Compose(transforms)

Composes several transforms together.

MakeOneHot()

Turn a categorical vector into a one-hot vector.

MakeOneHot2D(y_pixels, x_pixels)

Turn an array of continuous values into an array of one-hot 2D arrays.

MotionEnergy()

Compute motion energy across batch dimension.

SelectIdxs(idxs[, sample_name])

“Index-based subsampling of neural activity.

Threshold(threshold, bin_size)

Remove channels of neural activity whose mean value is below a threshold.

Transform()

Abstract base class for transforms.

ZScore()

z-score channel activity.

behavenet.data.utils Module

Utility functions for constructing inputs to data generators.

Functions

get_data_generator_inputs(hparams, sess_ids)

Helper function for generating signals, transforms and paths.

build_data_generator(hparams, sess_ids[, …])

Helper function to build data generator from hparams dict.

check_same_training_split(model_path, hparams)

Ensure data rng seed and trial splits are same for two models.

get_transforms_paths(data_type, hparams, sess_id)

Helper function for generating session-specific transforms and paths.

load_labels_like_latents(hparams, sess_ids, …)

Load labels from hdf5 in the same dictionary format that latents are saved.

get_region_list(hparams[, group_0, group_1])

Get brain regions and their indices into neural data.

behavenet.fitting

Model fitting documentation.

behavenet.fitting.eval Module

Utility functions for evaluating model fits.

Functions

export_latents(data_generator, model[, filename])

Export predicted latents using an already initialized data_generator and model.

export_predictions(data_generator, model[, …])

Export decoder predictions using an already initialized data_generator and model.

export_states(hparams, data_generator, model)

Export predicted latents using an already initialized data_generator and model.

export_train_plots(hparams, dtype[, …])

Export plot with MSE/LL as a function of training epochs.

get_reconstruction(model, inputs[, dataset, …])

Reconstruct an image from either image or latent inputs.

get_test_metric(hparams, model_version[, …])

Calculate a single R2 value across all test batches for a decoder.

behavenet.fitting.losses Module

Custom losses for PyTorch models.

Functions

mse(y_pred, y_true[, masks])

Compute mean square error (MSE) loss with masks.

gaussian_ll(y_pred, y_mean[, masks, std])

Compute multivariate Gaussian log-likelihood with a fixed diagonal noise covariance matrix.

gaussian_ll_to_mse(ll, n_dims[, …])

Convert a Gaussian log-likelihood term to MSE by removing constants and swapping variances.

kl_div_to_std_normal(mu, logvar)

Compute element-wise KL(q(z) || N(0, 1)) where q(z) is a normal parameterized by mu, logvar.

index_code_mi(z, mu, logvar)

Estimate index code mutual information in a batch.

total_correlation(z, mu, logvar)

Estimate total correlation in a batch.

dimension_wise_kl_to_std_normal(z, mu, logvar)

Estimate dimensionwise KL divergence to standard normal in a batch.

decomposed_kl(z, mu, logvar)

Decompose KL term in VAE loss.

subspace_overlap(A, B[, C])

Compute inner product between subspaces defined by matrices A and B.

triplet_loss(triplet_loss_obj, z, datasets)

Compute triplet loss to learn separated embedding space.

behavenet.fitting.training Module

Functions and classes for fitting PyTorch models with stochastic gradient descent.

Functions

fit(hparams, model, data_generator, exp[, …])

Fit pytorch models with stochastic gradient descent and early stopping.

Classes

Logger([n_datasets])

Base method for logging loss metrics.

EarlyStopping([patience, min_epochs, delta])

Stop training when a monitored quantity has stopped improving.

behavenet.fitting.utils Module

Utility functions for managing model paths and the hparams dict.

Functions

get_subdirs(path)

Get all first-level subdirectories in a given path (no recursion).

get_session_dir(hparams[, session_source])

Get session-level directory for saving model outputs from hparams dict.

get_expt_dir(hparams[, model_class, …])

Get output directories associated with a particular model class/type/testtube expt name.

read_session_info_from_csv(session_file)

Read csv file that contains lab/expt/animal/session info.

export_session_info_to_csv(session_dir, ids_list)

Export list of sessions to csv file.

contains_session(session_dir, session_id)

Determine if session defined by session_id dict is in the multi-session session_dir.

find_session_dirs(hparams)

Find all session dirs (single- and multi-session) that contain the session in hparams.

experiment_exists(hparams[, which_version])

Search testtube versions to find if experiment with the same hyperparameters has been fit.

get_model_params(hparams)

Returns dict containing all params considered essential for defining a model in that class.

export_hparams(hparams, exp)

Export hyperparameter dictionary.

get_lab_example(hparams, lab, expt)

Helper function to load data-specific hyperparameters and update hparams.

get_region_dir(hparams)

Return brain region string that combines region name and inclusion info.

create_tt_experiment(hparams)

Create test-tube experiment for logging training and storing models.

get_best_model_version(expt_dir[, measure, …])

Get best model version from a test tube experiment.

get_best_model_and_data(hparams[, Model, …])

Load the best model (and data) defined by hparams out of all available test-tube versions.

behavenet.models

Model documentation.

behavenet.models.base Module

Base models/modules in PyTorch.

Classes

BaseModule(*args, **kwargs)

Template for PyTorch modules.

BaseModel(*args, **kwargs)

Template for PyTorch models.

DiagLinear(features[, bias])

Applies a diagonal linear transformation to the incoming data: \(y = xD^T + b\)

CustomDataParallel(module[, device_ids, …])

Wrapper class for multi-gpu training.

behavenet.models.ae_model_architecture_generator Module

Functions

calculate_output_dim(input_dim, kernel, …)

Calculate output dimension of a layer/dimension based on input size, kernel size, etc.

draw_archs(batch_size, input_dim, n_ae_latents)

Generate multiple random autoencoder architectures with a fixed number of latents.

estimate_model_footprint(model, input_dim[, …])

Estimate model size to determine if it will fit on a single GPU.

get_decoding_conv_block(arch)

Build symmetric decoding block of convolutional autoencoder based on encoding block.

get_encoding_conv_block(arch, opts)

Build encoding block of convolutional autoencoder.

get_handcrafted_dims(arch[, symmetric])

Compute input/output dims as well as necessary padding for handcrafted architectures.

get_possible_arch(input_dim, n_ae_latents[, …])

Generate a random autoencoder architecture.

load_default_arch()

Load default convolutional AE architecture used in Whiteway et al 2021.

load_handcrafted_arch(input_dim, …[, …])

Load handcrafted autoencoder architecture from a json file.

load_handcrafted_arches(input_dim, …[, …])

Load handcrafted autoencoder architectures from a json file.

behavenet.models.aes Module

Autoencoder models implemented in PyTorch.

Functions

load_pretrained_ae(model, hparams)

Load pretrained weights into already constructed AE model.

Classes

ConvAEEncoder(hparams)

Convolutional encoder.

ConvAEDecoder(hparams)

Convolutional decoder.

LinearAEEncoder(n_latents, input_size)

Linear encoder.

LinearAEDecoder(n_latents, output_size[, …])

Linear decoder.

AE(hparams)

Base autoencoder class.

ConditionalAE(hparams)

Conditional autoencoder class.

AEMSP(hparams)

Autoencoder class with matrix subspace projection for disentangling the latent space.

behavenet.models.vaes Module

Variational autoencoder models implemented in PyTorch.

Functions

reparameterize(mu, logvar)

Sample from N(mu, var)

Classes

VAE(hparams)

Base variational autoencoder class.

ConditionalVAE(hparams)

Conditional variational autoencoder class.

BetaTCVAE(hparams)

Beta Total Correlation VAE class.

PSVAE(hparams)

Partitioned subspace variational autoencoder class.

MSPSVAE(hparams)

Partitioned subspace variational autoencoder class for multiple sessions.

ConvAEPSEncoder(hparams)

Convolutional encoder that separates label-related subspace.

ConvAEMSPSEncoder(hparams)

Convolutional encoder that separates label-related subspace.

behavenet.models.decoders Module

Encoding/decoding models implemented in PyTorch.

Classes

Decoder(hparams)

General wrapper class for encoding/decoding models.

MLP(hparams)

Feedforward neural network model.

LSTM(hparams)

LSTM neural network model.

ConvDecoder(hparams)

Decode images from predictors with a convolutional decoder.

behavenet.plotting

Plotting and video documentation.

behavenet.plotting Package

Utility functions shared across multiple plotting modules.

Functions

concat(ims[, axis])

Concatenate two channels along x or y direction (useful for data with multiple views).

get_crop(im, y_0, y_ext, x_0, x_ext)

Get crop of image, filling in borders with zeros.

load_latents(hparams, version[, dtype])

Load all latents as a single array.

load_metrics_csv_as_df(hparams, lab, expt, …)

Load metrics csv file and return as a pandas dataframe for easy plotting.

save_movie(save_file, ani[, frame_rate])

Save out matplotlib ArtistAnimation

behavenet.plotting.ae_utils Module

Plotting and video making functions for autoencoders.

Functions

make_ae_reconstruction_movie_wrapper(…[, …])

Produce movie with original video, reconstructed video, and residual.

make_reconstruction_movie(ims[, titles, …])

Produce movie with original video and reconstructed videos.

behavenet.plotting.cond_ae_utils Module

Functions

get_input_range(input_type, hparams[, …])

Helper function to compute input range for a variety of data types.

compute_range(values_list[, min_p, max_p])

Compute min and max of a list of numbers using percentiles.

get_labels_2d_for_trial(hparams, sess_ids[, …])

Return scaled labels (in pixel space) for a given trial.

get_model_input(data_generator, hparams, model)

Return images, latents, and labels for a given trial.

interpolate_2d(interp_type, model, ims_0, …)

Return reconstructed images created by interpolating through latent/label space.

interpolate_1d(interp_type, model, ims_0, …)

Return reconstructed images created by interpolating through latent/label space.

interpolate_point_path(interp_type, model, …)

Return reconstructed images created by interpolating through multiple points.

plot_2d_frame_array(ims_list[, markers, …])

Plot list of list of interpolated images output by interpolate_2d() in a 2d grid.

plot_1d_frame_array(ims_list[, markers, …])

Plot list of list of interpolated images output by interpolate_1d() in a 2d grid.

make_interpolated(ims, save_file[, markers, …])

Make a latent space interpolation movie.

make_interpolated_multipanel(ims, save_file)

Make a multi-panel latent space interpolation movie.

fit_classifier(model, data_generator[, …])

Fit classifier model from latent space to session id.

plot_psvae_training_curves(lab, expt, …[, …])

Create training plots for each term in the ps-vae objective function.

plot_hyperparameter_search_results(lab, …)

Create a variety of diagnostic plots to assess the ps-vae hyperparameters.

plot_label_reconstructions(lab, expt, …[, …])

Plot labels and their reconstructions from an ps-vae.

plot_latent_traversals(lab, expt, animal, …)

Plot video frames representing the traversal of individual dimensions of the latent space.

make_latent_traversal_movie(lab, expt, …)

Create a multi-panel movie with each panel showing traversals of an individual latent dim.

plot_mspsvae_training_curves(hparams, alpha, …)

Create training plots for each term in the ps-vae objective function.

plot_mspsvae_hyperparameter_search_results(…)

Create a variety of diagnostic plots to assess the msps-vae hyperparameters.

make_session_swap_movie(sess_ids, hparams, …)

Create a multipanel movie, each panel showing reconstruction with different session context.

behavenet.plotting.arhmm_utils Module

Plotting and video making functions for ARHMMs.

Functions

get_discrete_chunks(states[, include_edges])

Find occurences of each discrete state.

get_state_durations(latents, hmm[, …])

Calculate frame count for each state.

get_latent_arrays_by_dtype(data_generator[, …])

Collect data from data generator and put into dictionary with dtypes for keys.

get_model_latents_states(hparams, version[, …])

Return arhmm defined in hparams with associated latents and states.

make_syllable_movies_wrapper(hparams, save_file)

Present video clips of each individual syllable in separate panels.

make_syllable_movies(ims_orig, state_list, …)

Present video clips of each individual syllable in separate panels

real_vs_sampled_wrapper(output_type, …[, …])

Produce movie with (AE) reconstructed video and sampled video.

make_real_vs_sampled_movies(ims_recon, …)

Produce movie with (AE) reconstructed video and sampled video.

plot_real_vs_sampled(latents, latents_samp, …)

Plot real and sampled latents overlaying real and (potentially sampled) states.

plot_states_overlaid_with_latents(latents, …)

Plot states for a single trial overlaid with latents.

plot_state_transition_matrix(model[, deridge])

Plot Markov transition matrix for arhmm.

plot_dynamics_matrices(model[, deridge])

Plot autoregressive dynamics matrices for arhmm.

plot_obs_biases(model)

Plot observation bias vectors for arhmm.

plot_obs_covariance_matrices(model)

Plot observation covariance matrices for arhmm.

behavenet.plotting.decoder_utils Module

Plotting functions for decoders.

Functions

get_r2s_by_trial(hparams, model_types)

For a given session, load R^2 metrics from all decoders defined by hparams.

get_best_models(metrics_df)

Find best decoder over l2 regularization and learning rate.

get_r2s_across_trials(hparams, best_models_df)

Calculate R^2 across all test trials (rather than on a trial-by-trial basis)

make_neural_reconstruction_movie_wrapper(…)

Produce movie with original video, ae reconstructed video, and neural reconstructed video.

make_neural_reconstruction_movie(ims_orig, …)

Produce movie with original video, ae reconstructed video, and neural reconstructed video.

plot_neural_reconstruction_traces_wrapper(hparams)

Plot ae latents and their neural reconstructions.

plot_neural_reconstruction_traces(traces_ae, …)

Plot ae latents and their neural reconstructions.

Indices and tables