Utils

Functions for the common machine learning utilities plus diffusion related methods.

Loss function

Loss


def Loss(
    RF:int, c:int, use_pad:bool=False, beta:float=0.0, reduction:str='sum', log:bool=False
):

β-VAE loss function for a reconstruction term that is the negative log-likelihood of a mixture of Gaussians

Gaussian mixture

source

mix_gaussian_loss


def mix_gaussian_loss(
    y_hat:Tensor, # Predicted output (B x C x T)
    y:Tensor, # Target (B x T x 1)
    log_scale_min:float=-12.0, # Log scale minimum value
    reduction:str='sum'
)->Tensor: # Loss

Mixture of continuous Gaussian distributions loss. Note that it is assumed that input is scaled to [-1, 1].

Kullback-Leibler Divergence

source

kl_divergence


def kl_divergence(
    mu:Tensor, logvar:Tensor, reduction:str='meansum'
):

Divergence term of the VAE loss function.

Beta-TCVAE Loss

source

get_log_pz_qz_prodzi_qzCx


def get_log_pz_qz_prodzi_qzCx(
    latent_sample, latent_dist, n_data, is_mss:bool=True
):

source

log_importance_weight_matrix


def log_importance_weight_matrix(
    batch_size:int, dataset_size:int
):

Calculates a log importance weight matrix.

source

matrix_log_density_gaussian


def matrix_log_density_gaussian(
    x, mu, logvar
):

Calculates log density of a Gaussian for all batch pairs.

source

log_density_gaussian


def log_density_gaussian(
    x, mu, logvar
):

Calculates log density of a Gaussian.

source

BtcvaeLoss


def BtcvaeLoss(
    RF, c, use_pad:bool=False, alpha:float=0.0001, beta:float=3.0, gamma:float=0.1, reduction:str='sum',
    is_mss:bool=True, log:bool=True
):

Beta-TCVAE loss with modular gamma, KL, and MoG log-likelihood.

Metrics

source

GaussianMixtureMetric


def GaussianMixtureMetric(
    RF:int, c:int, use_pad:bool=False, func:function=mix_gaussian_loss, reduction:str='mean'
):

Metric to log the Gaussian mixture loss

source

KLDMetric


def KLDMetric(
    c:int
):

Metric to log the Kullback-Leibler divergence term

source