Function main

quantization/eval_mobilenetv2.js:73–141

()

}

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

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

  const imageH = model.inputs[0].shape[2];
  const imageW = model.inputs[0].shape[2];

  // Load the images into a tensor.
  const dirContent = fs.readdirSync(args.imageDir);
  dirContent.sort();
  const numImages = dirContent.length;
  console.log(`Reading ${numImages} images...`);
  const progressBar = new ProgressBar('[:bar]', {
    total: numImages,
    width: 80,
    head: '>'
  });
  const imageTensors = [];
  const truthLabels = [];
  for (const fileName of dirContent) {
    const truthLabel = fileName.split('.')[0].split('_')[2];
    truthLabels.push(truthLabel);
    const imageFilePath = path.join(args.imageDir, fileName);
    const imageTensor =
        await readImageTensorFromFile(imageFilePath, imageH, imageW);
    imageTensors.push(imageTensor);
    progressBar.tick();
  }

  const stackedImageTensor = tf.concat(imageTensors, 0);
  console.log('Calling model.predict()...');
  const t0 = new Date().getTime();
  const {top1Indices, top5Indices} = tf.tidy(() => {
    const probs = model.predict(stackedImageTensor, {batchSize: 64});
    return {
      top1Indices: probs.argMax(-1).arraySync(),
      top5Indices: probs.topk(5).indices.arraySync()
    };
  });
  console.log(`model.predict() took ${(new Date().getTime() - t0).toFixed(2)} ms`);

  let numCorrectTop1 = 0;
  let numCorrectTop5 = 0;
  top1Indices.forEach((top1Index, i) => {
    const truthLabel = truthLabels[i];
    if (IMAGENET_CLASSES[top1Index].indexOf(truthLabel) !== -1) {
      numCorrectTop1++;
    }
    for (let k = 0; k < 5; ++k) {
      if (IMAGENET_CLASSES[top5Indices[i][k]].indexOf(truthLabel) !== -1) {
        numCorrectTop5++;