Function main

quantization/train_housing.js:69–105

()

}

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

  const {count, featureMeans, featureStddevs, labelMean, labelStddev} =
      await getDatasetStats();
  const {trainXs, trainYs, valXs, valYs, evalXs, evalYs} =
      await getNormalizedDatasets(
          count, featureMeans, featureStddevs, labelMean, labelStddev,
          args.validationSplit, args.evaluationSplit);

  const model = createModel();
  model.summary();

  await model.fit(trainXs, trainYs,  {
    epochs: args.epochs,
    batchSize: args.batchSize,
    validationData: [valXs, valYs]
  });

  const evalOutput = model.evaluate(evalXs, evalYs);
  console.log(
      `\nEvaluation result:\n` +
      `  Loss = ${evalOutput.dataSync()[0].toFixed(6)}`);

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

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