Function train

fashion-mnist-vae/train.js:50–112

* Train the auto encoder * * @param {number[][]} images Flattened images for VAE training. * @param {object} vaeOpts Options for the VAE model, including the following * fields: * - originaDim {number} Length of the input flattened image. * - intermediateDim {number} Number of units of t

(images, vaeOpts, savePath, logDir)

 *   so that the training process can be monitored using TensorBoard.
 */
  async function train(images, vaeOpts, savePath, logDir) {
  const encoderModel = encoder(vaeOpts);
  const decoderModel = decoder(vaeOpts);
  const vaeModel = vae(encoderModel, decoderModel);

  let summaryWriter;
  if (logDir != null) {
    console.log(`Logging loss values to ${logDir}.`);
    console.log(
        `Use the following command to start the tensorboard backend server:`);
    console.log(`  tensorboard --logdir ${logDir}`);
    summaryWriter = tf.node.summaryFileWriter(logDir);
  }

  console.log('\n** Train Model **\n');

  // Because we use a custom loss function, we will use optimizer.minimize
  // instead of the more typical model.fit. We thus need to define an optimizer
  // and manage batching the data ourselves.

  // Create the optimizer
  const optimizer = tf.train.adam();

  // Group the data into batches.
  const batches = _.chunk(images, batchSize);

  // Run the train loop.
  let step = 0;
  for (let i = 0; i < epochs; i++) {
    console.log(`\nEpoch #${i + 1} of ${epochs}\n`);
    for (let j = 0; j < batches.length; j++) {
      const currentBatchSize = batches[j].length
      const batchedImages = batchImages(batches[j]);

      const reshaped =
          batchedImages.reshape([currentBatchSize, vaeOpts.originalDim]);

      // This is the model optimization step. We make a prediction
      // compute loss and return it so that optimizer.minimize can
      // adjust the weights of the model.
      optimizer.minimize(() => {
        const outputs = vaeModel.apply(reshaped);
        const loss = vaeLoss(reshaped, outputs, vaeOpts);
        process.stdout.write('.');
        if (j % 50 === 0) {
          console.log('\nLoss:', loss.dataSync()[0]);
        }
        if (summaryWriter != null) {
          summaryWriter.scalar('loss', loss, step++);
        }

        return loss;
      });
      tf.dispose([batchedImages, reshaped]);
    }
    console.log('');
    // Generate a preview after each epoch
    await generate(decoderModel, vaeOpts.latentDim);