Adding extra loss to constrain the latent variable in VAE

[utkarshp1161]

Hi Maxim,

Thanks again for the great package.

I wish to accomplish the following:
I have a version of VAE in Keras with the following loss terms (following our work https://arxiv.org/abs/2402.02198):

1. KL_divergence loss
2. Custom_loss: Essentially, each x has an associated y. I wish to constrain one of the latent variables to this y value.
3. Reconstruction loss

reconstruction_loss = binary_crossentropy(K.flatten(x), K.flatten(reconstructed))
reconstruction_loss *= self.image_size

# KL divergence
kl_loss = 1 + z_log_var - K.square(z_mean) - K.exp(z_log_var)
kl_loss = -0.5 * K.sum(kl_loss, axis=-1)

# Custom loss
batch_size = tf.shape(x)[0]
y_true = tf.cast(y[:batch_size], dtype='float32')
custom_loss = K.mean(K.square(z[:, 0] - y_true), axis=-1) * 10

# Total loss
total_loss = K.mean(reconstruction_loss + kl_loss + custom_loss)

How to add this extra loss (custom loss) in the PyroVED iVAE class? Do you have any suggestions on where to add this loss term?

Thank you!

[ziatdinovmax]

You will need to define a custom loss in pyro following these guidelines: https://pyro.ai/examples/custom_objectives.html