Function run

abalone-node/train.js:33–55

* Train a model with dataset, then save the model to a local folder.

(epochs, batchSize, savePath)

 * Train a model with dataset, then save the model to a local folder.
 */
async function run(epochs, batchSize, savePath) {
  const datasetObj = await createDataset('file://' + csvPath);
  const model = createModel([datasetObj.numOfColumns]);
  // The dataset has 4177 rows. Split them into 2 groups, one for training and
  // one for validation. Take about 3500 rows as train dataset, and the rest as
  // validation dataset.
  const trainBatches = Math.floor(3500 / batchSize);
  const dataset = datasetObj.dataset.shuffle(1000).batch(batchSize);
  const trainDataset = dataset.take(trainBatches);
  const validationDataset = dataset.skip(trainBatches);

  await model.fitDataset(
      trainDataset, {epochs: epochs, validationData: validationDataset});

  await model.save(savePath);

  const loadedModel = await tf.loadLayersModel(savePath + '/model.json');
  const result = loadedModel.predict(
      tf.tensor2d([[0, 0.625, 0.495, 0.165, 1.262, 0.507, 0.318, 0.39]]));
  console.log(
      'The actual test abalone age is 10, the inference result from the model is ' +
      result.dataSync());
}

const parser = new argparse.ArgumentParser(
    {description: 'TensorFlow.js-Node Abalone Example.', addHelp: true});