Function trainDiscriminatorOneStep

mnist-acgan/gan.js:298–329

* Train the discriminator for one step. * * In this step, only the weights of the discriminator are updated. The * generator is not involved. * * The following steps are involved: * * - Slice the training features and to get batch of real data. * - Generate a random latent-space vector a

(
    xTrain, yTrain, batchStart, batchSize, latentSize, generator,
    discriminator)

 * @returns {number[]} The loss values from the one-step training as numbers.
 */
async function trainDiscriminatorOneStep(
    xTrain, yTrain, batchStart, batchSize, latentSize, generator,
    discriminator) {
  // TODO(cais): Remove tidy() once the current memory leak issue in tfjs-node
  //   and tfjs-node-gpu is fixed.
  const [x, y, auxY] = tf.tidy(() => {
    const imageBatch = xTrain.slice(batchStart, batchSize);
    const labelBatch = yTrain.slice(batchStart, batchSize).asType('float32');

    // Latent vectors.
    let zVectors = tf.randomUniform([batchSize, latentSize], -1, 1);
    let sampledLabels =
        tf.randomUniform([batchSize, 1], 0, NUM_CLASSES, 'int32')
            .asType('float32');

    const generatedImages =
        generator.predict([zVectors, sampledLabels], {batchSize: batchSize});

    const x = tf.concat([imageBatch, generatedImages], 0);

    const y = tf.tidy(
        () => tf.concat(
            [tf.ones([batchSize, 1]).mul(SOFT_ONE), tf.zeros([batchSize, 1])]));

    const auxY = tf.concat([labelBatch, sampledLabels], 0);
    return [x, y, auxY];
  });

  const losses = await discriminator.trainOnBatch(x, [y, auxY]);
  tf.dispose([x, y, auxY]);
  return losses;
}

/**
 * Train the combined ACGAN for one step.