Function vaeLoss

fashion-mnist-vae/model.js:179–203

* The custom loss function for VAE. * * @param {tf.tensor} inputs the encoder inputs a batched image tensor * @param {[tf.tensor]} outputs the vae outputs, [decoderOutput, * ...encoderOutputs] * @param {number} vaeOpts.originalDim number of dimensions in the original data

(inputs, outputs)

 * @param {number} vaeOpts.originalDim number of dimensions in the original data
 */
function vaeLoss(inputs, outputs) {
  return tf.tidy(() => {
    const originalDim = inputs.shape[1];
    const decoderOutput = outputs[0];
    const zMean = outputs[1];
    const zLogVar = outputs[2];

    // First we compute a 'reconstruction loss' terms. The goal of minimizing
    // this term is to make the model outputs match the input data.
    const reconstructionLoss =
        tf.losses.meanSquaredError(inputs, decoderOutput).mul(originalDim);

    // binaryCrossEntropy can be used as an alternative loss function
    // const reconstructionLoss =
    //  tf.metrics.binaryCrossentropy(inputs, decoderOutput).mul(originalDim);

    // Next we compute the KL-divergence between zLogVar and zMean, minimizing
    // this term aims to make the distribution of latent variable more normally
    // distributed around the center of the latent space.
    let klLoss = zLogVar.add(1).sub(zMean.square()).sub(zLogVar.exp());
    klLoss = klLoss.sum(-1).mul(-0.5);

    return reconstructionLoss.add(klLoss).mean();
  });
}

module.exports = {
  vae,