PSVAE

class behavenet.models.vaes.PSVAE(hparams)

Bases: AE

Partitioned subspace variational autoencoder class.

This class constructs a VAE that…

Methods Summary

build_model()	Construct the model using hparams.
forward(x[, dataset, use_mean])	Process input data.
get_inverse_transformed_latents(inputs[, ...])	Return latents after they have been transformed using the diagonal mapping D.
get_predicted_labels(x[, dataset, use_mean])	Process input data to get predicted labels.
get_transformed_latents(inputs[, dataset, ...])	Return latents after supervised subspace has been transformed to original label space.
loss(data[, dataset, accumulate_grad, ...])	Calculate modified ELBO loss for PSVAE.

Methods Documentation

build_model()

Construct the model using hparams.

forward(x, dataset=None, use_mean=False, **kwargs)

Process input data.

Parameters:

• x (torch.Tensor object) – input data of shape (n_frames, n_channels, y_pix, x_pix)

• dataset (int) – used with session-specific io layers

• use_mean (bool) – True to skip sampling step

Returns:

• x_hat (torch.Tensor): output of shape (n_frames, n_channels, y_pix, x_pix)

• z (torch.Tensor): sampled latent variable of shape (n_frames, n_latents)

• mu (torch.Tensor): mean paramter of shape (n_frames, n_latents)

• logvar (torch.Tensor): logvar paramter of shape (n_frames, n_latents)

• y_hat (torch.Tensor): output of shape (n_frames, n_labels)

Return type:

tuple

get_inverse_transformed_latents(inputs, dataset=None, as_numpy=True)

Return latents after they have been transformed using the diagonal mapping D.

Parameters:

• inputs (torch.Tensor object) –

  • image tensor of shape (n_frames, n_channels, y_pix, x_pix)

  • latents tensor of shape (n_frames, n_ae_latents) where the first n_labels entries are assumed to be labels in the original pixel space

• dataset (int, optional) – used with session-specific io layers

• as_numpy (bool, optional) – True to return as numpy array, False to return as torch Tensor

Returns:

array of latents in transformed latent space of shape (n_frames, n_latents)

Return type:

np.ndarray or torch.Tensor object

get_predicted_labels(x, dataset=None, use_mean=True)

Process input data to get predicted labels.

Parameters:

• x (torch.Tensor object) – input data of shape (n_frames, n_channels, y_pix, x_pix)

• dataset (int) – used with session-specific io layers

• use_mean (bool) – True to skip sampling step

Returns:

output of shape (n_frames, n_labels)

Return type:

torch.Tensor

get_transformed_latents(inputs, dataset=None, as_numpy=True)

Return latents after supervised subspace has been transformed to original label space.

Parameters:

• inputs (torch.Tensor object) –

  • image tensor of shape (n_frames, n_channels, y_pix, x_pix)

  • latents tensor of shape (n_frames, n_ae_latents)

• dataset (int, optional) – used with session-specific io layers

• as_numpy (bool, optional) – True to return as numpy array, False to return as torch Tensor

Returns:

array of latents in transformed latent space of shape (n_frames, n_latents)

Return type:

np.ndarray or torch.Tensor object

loss(data, dataset=0, accumulate_grad=True, chunk_size=200)

Calculate modified ELBO loss for PSVAE.

The batch is split into chunks if larger than a hard-coded chunk_size to keep memory requirements low; gradients are accumulated across all chunks before a gradient step is taken.

Parameters:

• data (dict) – batch of data; keys should include ‘images’ and ‘masks’, if necessary

• dataset (int, optional) – used for session-specific io layers

• accumulate_grad (bool, optional) – accumulate gradient for training step

• chunk_size (int, optional) – batch is split into chunks of this size to keep memory requirements low

Returns:

• ‘loss’ (float): full elbo

• ’loss_ll’ (float): log-likelihood portion of elbo

• ’loss_kl’ (float): kl portion of elbo

• ’loss_mse’ (float): mse (without gaussian constants)

• ’beta’ (float): weight in front of kl term

Return type:

dict