Add a new Node.js example: abalone (#328)

* add new example

* switch to cpu

* address comments

* address comments

* add test data

* update

* remove nit

* update

Author: Kangyi Zhang

README.md

@@ -19,6 +19,17 @@ to another project.
     <th>Inference</th>
     <th>API type</th>
     <th>Save-load operations</th>
+  <tr>
+    <td><a href="./abalone-node">abalone-node</a></td>
+    <td></td>
+    <td>Numeric</td>
+    <td>Loading data from local file and training in Node.js</td>
+    <td>Multilayer perceptron</td>
+    <td>Node.js</td>
+    <td>Node.js</td>
+    <td>Layers</td>
+    <td>Saving to filesystem and loading in Node.js</td>
+  </tr>
   <tr>
     <td><a href="./addition-rnn">addition-rnn</a></td>
     <td><a href="https://storage.googleapis.com/tfjs-examples/addition-rnn/dist/index.html">🔗</a></td>

abalone-node/README.md

@@ -0,0 +1,24 @@
+# TensorFlow.js Example: Abalone Age
+
+This example shows how to predicting the age of abalone from physical measurements under Node.js
+
+The data set available at [UCI Machine Learning Repository](https://archive.ics.uci.edu/ml/datasets/Abalone).
+
+This example shows how to
+* load a `Dataset` from a local csv file
+* prepare the Dataset for training
+* create a `tf.LayersModel` from scratch
+* train the model through `model.fitDataset()`
+* save the trained model to a local folder.
+
+To launch the demo, do
+
+```sh
+yarn
+yarn train
+```
+
+By default, the training uses tfjs-node, which runs on the CPU.
+If you have a CUDA-enabled GPU and have the CUDA and CuDNN libraries
+set up properly on your system, you can run the training on the GPU
+by replacing the tfjs-node package with tfjs-node-gpu.

abalone-node/data.js

@@ -0,0 +1,55 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+const tf = require('@tensorflow/tfjs-node');
+
+/**
+ * Load a local csv file and prepare the data for training. Data source:
+ * https://archive.ics.uci.edu/ml/datasets/Abalone
+ *
+ * @param {string} csvPath The path to csv file.
+ * @returns {tf.data.Dataset} The loaded and prepared Dataset.
+ */
+async function createDataset(csvPath) {
+  const dataset = tf.data.csv(
+      csvPath, {hasHeader: true, columnConfigs: {'rings': {isLabel: true}}});
+  const numOfColumns = (await dataset.columnNames()).length - 1;
+  // Convert features and labels.
+  return {
+    dataset: dataset.map(row => {
+      const rawFeatures = row['xs'];
+      const rawLabel = row['ys'];
+      const convertedFeatures = Object.keys(rawFeatures).map(key => {
+        switch (rawFeatures[key]) {
+          case 'F':
+            return 0;
+          case 'M':
+            return 1;
+          case 'I':
+            return 2;
+          default:
+            return Number(rawFeatures[key]);
+        }
+      });
+      const convertedLabel = [rawLabel['rings']];
+      return {xs: convertedFeatures, ys: convertedLabel};
+    }),
+    numOfColumns: numOfColumns
+  };
+}
+
+module.exports = createDataset;

abalone-node/data_test.js

@@ -0,0 +1,42 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+const tf = require('@tensorflow/tfjs-node');
+const createDataset = require('./data');
+
+describe('Dataset', () => {
+  it('Created dataset and numOfColumns', async () => {
+    const csvPath = 'file://./test_data.csv';
+    const datasetObj = await createDataset(csvPath);
+    const dataset = datasetObj.dataset;
+    const row = await dataset.take(1).toArray();
+    const numOfColumns = datasetObj.numOfColumns;
+    expect(numOfColumns).toBe(8);
+    const features = row[0].xs;
+    const label = row[0].ys;
+    expect(features.length).toBe(8);
+    expect(features[0] === 0 || features[0] === 1 || features[0] === 2)
+        .toBeTruthy();
+    for (let i = 1; i < 8; i++) {
+      expect(features[i]).toBeLessThan(1);
+      expect(features[i]).toBeGreaterThan(0);
+    }
+    expect(Number.isInteger(label[0])).toBeTruthy();
+    expect(label[0]).toBeGreaterThanOrEqual(2);
+    expect(label[0]).toBeLessThanOrEqual(16);
+  });
+});

abalone-node/model.js

@@ -0,0 +1,44 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+const tf = require('@tensorflow/tfjs-node');
+
+/**
+ * Builds and returns Multi Layer Perceptron Regression Model.
+ *
+ * @param {number} inputShape The input shape of the model.
+ * @returns {tf.Sequential} The multi layer perceptron regression mode  l.
+ */
+function createModel(inputShape) {
+  const model = tf.sequential();
+  model.add(tf.layers.dense({
+    inputShape: inputShape,
+    activation: 'sigmoid',
+    units: 50,
+  }));
+  model.add(tf.layers.dense({
+    activation: 'sigmoid',
+    units: 50,
+  }));
+  model.add(tf.layers.dense({
+    units: 1,
+  }));
+  model.compile({optimizer: tf.train.sgd(0.01), loss: 'meanSquaredError'});
+  return model;
+}
+
+module.exports = createModel;

abalone-node/model_test.js

@@ -0,0 +1,66 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+const tf = require('@tensorflow/tfjs-node');
+const shelljs = require('shelljs');
+const tmp = require('tmp');
+const fs = require('fs');
+const createModel = require('./model');
+
+let tempDir;
+
+describe('Model', () => {
+  beforeEach(() => {
+    tempDir = tmp.dirSync();
+  });
+
+  afterEach(() => {
+    if (fs.existsSync(tempDir)) {
+      shelljs.rm('-rf', tempDir);
+    }
+  });
+
+  it('Created model can train', async () => {
+    const inputLength = 6;
+    const outputLength = 1;
+    const model = createModel([inputLength]);
+    expect(model.inputs.length).toEqual(1);
+    expect(model.inputs[0].shape).toEqual([null, inputLength]);
+    expect(model.outputs.length).toEqual(1);
+    expect(model.outputs[0].shape).toEqual([null, outputLength]);
+
+    const numExamples = 3;
+    const inputFeature = tf.ones([numExamples, inputLength]);
+    const inputLabel = tf.ones([numExamples, outputLength]);
+    const history = await model.fit(inputFeature, inputLabel, {epochs: 2});
+    expect(history.history.loss.length).toEqual(2);
+  });
+
+  it('Model save-load roundtrip', async () => {
+    const inputLength = 6;
+    const model = createModel([inputLength]);
+    const numExamples = 3;
+    const feature = tf.ones([numExamples, inputLength]);
+    const y = model.predict(feature);
+
+    await model.save(`file://${tempDir.name}`);
+    const modelPrime =
+        await tf.loadLayersModel(`file://${tempDir.name}/model.json`);
+    const yPrime = modelPrime.predict([feature]);
+    tf.test_util.expectArraysClose(yPrime, y);
+  });
+});

abalone-node/package.json

@@ -0,0 +1,21 @@
+{
+    "name": "tfjs-abalone-node",
+    "version": "1.0.0",
+    "description": "",
+    "main": "index.js",
+    "license": "Apache-2.0",
+    "scripts": {
+        "train": "node train.js",
+        "test": "node run_tests.js"
+    },
+    "dependencies": {
+        "@tensorflow/tfjs-node": "1.2.8",
+        "argparse": "^1.0.10"
+    },
+    "devDependencies": {
+        "jasmine": "^3.2.0",
+        "jasmine-core": "^3.2.1",
+        "shelljs": "^0.8.3",
+        "tmp": "^0.0.33"
+    }
+}

abalone-node/run_tests.js

@@ -0,0 +1,21 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+const jasmine_util = require('@tensorflow/tfjs-core/dist/jasmine_util');
+const runTests = require('../test_util').runTests;
+
+runTests(jasmine_util, ['./*test.js']);

abalone-node/test_data.csv

@@ -0,0 +1,4 @@
+sex,length,diameter,height,weight.w,weight.s,weight.v,weight.sh,rings
+M,0.455,0.365,0.095,0.514,0.2245,0.101,0.15,15
+M,0.35,0.265,0.09,0.2255,0.0995,0.0485,0.07,7
+F,0.53,0.42,0.135,0.677,0.2565,0.1415,0.21,9

abalone-node/train.js

@@ -0,0 +1,80 @@
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+const tf = require('@tensorflow/tfjs-node');
+const argparse = require('argparse');
+const https = require('https');
+const fs = require('fs');
+const createModel = require('./model');
+const createDataset = require('./data');
+
+
+const csvUrl =
+    'https://storage.googleapis.com/tfjs-examples/abalone-node/abalone.csv';
+const csvPath = './abalone.csv';
+
+/**
+ * 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});
+parser.addArgument('--epochs', {
+  type: 'int',
+  defaultValue: 100,
+  help: 'Number of epochs to train the model for.'
+});
+parser.addArgument('--batch_size', {
+  type: 'int',
+  defaultValue: 500,
+  help: 'Batch size to be used during model training.'
+})
+parser.addArgument(
+    '--savePath',
+    {type: 'string', defaultValue: 'file://trainedModel', help: 'Path.'})
+const args = parser.parseArgs();
+
+
+const file = fs.createWriteStream(csvPath);
+https.get(csvUrl, function(response) {
+  response.pipe(file).on('close', async () => {
+    run(args.epochs, args.batch_size, args.savePath);
+  });
+});