Variational Autoencoder - custom loss

Hi all, could anyone please advise on how to express a custom loss function while training a VAE in Numpyro?

For instance, if in this example model defines the reconstruction_loss and the guide defines the KL_loss, how to express the total beta-loss (reconstruction_loss + beta* KL_loss)?

you can use the scale handler to scale terms in the model/guide:

def model(..., beta=0.5):
    with numpyro.handlers.scale(scale=beta):
        z = numpyro.sample("z", ...)
    return numpyro.sample("obs", dist.Bernoulli(img_loc), obs=batch)

def guide(..., beta=0,5):
    with numpyro.handlers.scale(scale=beta):
        z = numpyro.sample("z", dist.Normal(z_loc, z_std))
1 Like

@martinjankowiak thank you for the previous answer!

What I really need to do is to decompose the reconstruction loss into two components. For example, in the published example it would go along the lines of

v1 = img_loc[1:n1]
v2 = img_loc[n1:n2]

Then I would like to calculate the total VAE loss as w_1*L_1 + w_2*L_2 + KL, where L_1 and L_2 are log-likelihoods of v1 and v2, respectively, and w_1 and w_2 are some weights (fixed numbers).

How can this be implemented in Numpyro?

if i understand what you want you can still use the scale handler for this. the scale does not need to be a scalar. it can be e.g. a vector that is inside a plate that scales the likelihood for each data point differently

1 Like