Function trainModelUsingFitDataset

data-generator/index.js:135–173

* Trains a the provided model on the provided dataset using model.fitDataset. * Schedules a callback at the end of every epoch to update the UI with * graphs showing loss and accuracy, as well as training speed and the current * prediction for the manually entered hand. * @param {tf.Model} model

(model, dataset)

 * @param {tf.data.Dataset} dataset
 */
async function trainModelUsingFitDataset(model, dataset) {
  const trainLogs = [];
  const beginMs = performance.now();
  const fitDatasetArgs = {
    batchesPerEpoch: ui.getBatchesPerEpoch(),
    epochs: ui.getEpochsToTrain(),
    validationData: dataset,
    validationBatches: 10,
    callbacks: {
      onEpochEnd: async (epoch, logs) => {
        // Plot the loss and accuracy values at the end of every training epoch.
        const secPerEpoch =
            (performance.now() - beginMs) / (1000 * (epoch + 1));
        ui.displayTrainLogMessage(
            `Training model... Approximately ` +
            `${secPerEpoch.toFixed(4)} seconds per epoch`);
        trainLogs.push(logs);
        tfvis.show.history(
            ui.lossContainerElement, trainLogs, ['loss', 'val_loss'])
        tfvis.show.history(
            ui.accuracyContainerElement, trainLogs, ['acc', 'val_acc'],
            {zoomToFitAccuracy: true})
        ui.displayNumSimulationsSoFar();
        // Update the prediction.
        predictHandler();
        // Stop the training if stop requested.
        if (STOP_REQUESTED) {
          model.stopTraining = true;
        }
      },
    }
  };
  ui.disableTrainButton();
  ui.enableStopButton();
  ui.enablePredictButton();
  await model.fitDataset(dataset, fitDatasetArgs);
  ui.enableTrainButton();
  ui.disableStopButton();
}

/**
 * Constructs a new model and trains it on a dataset pipeline built off of