Function main

quantization/eval_mnist.js:54–86

()

}

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

  let dataset;
  if (args.dataset === 'fashion-mnist') {
    dataset = new FashionMnistDataset();
  } else if (args.dataset === 'mnist') {
    dataset = new MnistDataset();
  } else {
    throw new Error(`Unrecognized dataset name: ${args.dataset}`);
  }
  await dataset.loadData();
  const {images: testImages, labels: testLabels} = dataset.getTestData();

  console.log(`Loading model from ${args.modelSavePath}...`);
  const model = await tf.loadLayersModel(`file://${args.modelSavePath}`);
  compileModel(model);

  console.log(`Performing evaluation...`);
  const t0 = tf.util.now();
  const evalOutput = model.evaluate(testImages, testLabels);
  const t1 = tf.util.now();
  console.log(`\nEvaluation took ${(t1 - t0).toFixed(2)} ms.`);
  console.log(
      `\nEvaluation result:\n` +
      `  Loss = ${evalOutput[0].dataSync()[0].toFixed(6)}; `+
      `Accuracy = ${evalOutput[1].dataSync()[0].toFixed(6)}`);
}

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