Function main

quantization/train_mnist.js:68–115

()

}

async function main() {
  const args = parseArgs();
  if (args.gpu) {
    tf = require('@tensorflow/tfjs-node-gpu');
  } else {
    tf = require('@tensorflow/tfjs-node');
  }

  let dataset;
  let model;
  if (args.dataset === 'fashion-mnist') {
    dataset = new FashionMnistDataset();
    model = createFashionMnistModel();
  } else if (args.dataset === 'mnist') {
    dataset = new MnistDataset();
    model = createMnistModel();
  } else {
    throw new Error(`Unrecognized dataset name: ${args.dataset}`);
  }
  await dataset.loadData();
  const {images: trainImages, labels: trainLabels} = dataset.getTrainData();

  model.summary();

  await model.fit(trainImages, trainLabels, {
    epochs: args.epochs,
    batchSize: args.batchSize,
    validationSplit: args.validationSplit,
    callbacks: tf.callbacks.earlyStopping({patience: 20})
  });

  const {images: testImages, labels: testLabels} = dataset.getTestData();
  const evalOutput = model.evaluate(testImages, testLabels);

  console.log(
      `\nEvaluation result:\n` +
      `  Loss = ${evalOutput[0].dataSync()[0].toFixed(6)}; `+
      `Accuracy = ${evalOutput[1].dataSync()[0].toFixed(6)}`);

  const modelSavePath = path.join(args.modelSavePath, args.dataset, 'original');
  if (modelSavePath != null) {
    if (!fs.existsSync(path.dirname(modelSavePath))) {
      shelljs.mkdir('-p', path.dirname(modelSavePath));
    }
    await model.save(`file://${modelSavePath}`);
    console.log(`Saved model to path: ${modelSavePath}`);
  }
}

if (require.main === module) {
  main();