Function train

mnist-core/model.js:78–90

(data, log)

// Train the model.
export async function train(data, log) {
  const returnCost = true;

  for (let i = 0; i < TRAIN_STEPS; i++) {
    const cost = optimizer.minimize(() => {
      const batch = data.nextTrainBatch(BATCH_SIZE);
      return loss(batch.labels, model(batch.xs));
    }, returnCost);

    log(cost.dataSync(), i);
    await tf.nextFrame();
  }
}

// Predict the digit number from a batch of input images.
export function predict(x) {