[jena-weather] Fix baseline; add unit tests; upgrade to tfjs 1.1.0 (tensorflow#271)

- Fix a crash in the `--modelType baseline` command
- Add unit tests
- Upgrade to tfjs and tfjs-node/tfjs-node-gpu 1.1.0
- Add early stopping with a default patience of 2

Author: caisq

jena-weather/data_test.js

@@ -0,0 +1,44 @@
+/**
+ * @license
+ * Copyright 2018 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.
+ * =============================================================================
+ */
+
+import {JenaWeatherData} from "./data";
+
+global.fetch = require('node-fetch');
+
+describe('JenaWeatherData', () => {
+  it('construct, load and basic public methods', async () => {
+    const dataset = new JenaWeatherData();
+    await dataset.load();
+
+    expect(dataset.getDataColumnNames().length).toEqual(14);
+    expect(dataset.getDataColumnNames()[0]).toEqual('p (mbar)');
+    expect(dataset.getDataColumnNames()[1]).toEqual('T (degC)');
+
+    expect(new Date(dataset.getTime(0)).getTime()).toBeGreaterThan(0);
+
+    const columnData = dataset.getColumnData('T (degC)', false, true, 0, 30, 3);
+    expect(columnData.length).toEqual(10);
+
+    const func = dataset.getNextBatchFunction(
+        true, 1000, 100, 32, 10, 0, 10000, true, false);
+    for (let i = 0; i < 2; ++i) {
+      const item = func.next();
+      expect(item.done).toEqual(false);
+      expect(item.value.xs.shape).toEqual([32, 100, 14]);
+    }
+  });
+});

jena-weather/models.js

@@ -16,7 +16,8 @@
  */
 
 /**
- * Creating and training `tf.Model`s for the temperature prediction problem.
+ * Creating and training `tf.LayersModel`s for the temperature prediction
+ * problem.
  *
  * This file is used to create models for both
  * - the browser: see [index.js](./index.js), and
@@ -51,23 +52,24 @@ const VAL_MAX_ROW = 300000;
 export async function getBaselineMeanAbsoluteError(
     jenaWeatherData, normalize, includeDateTime, lookBack, step, delay) {
   const batchSize = 128;
-  const nextBatchFn = jenaWeatherData.getNextBatchFunction(
-      false, lookBack, delay, batchSize, step, VAL_MIN_ROW, VAL_MAX_ROW,
-      normalize, includeDateTime);
-  const dataset = tf.data.generator(nextBatchFn);
+  const dataset = tf.data.generator(
+    () => jenaWeatherData.getNextBatchFunction(
+        false, lookBack, delay, batchSize, step, VAL_MIN_ROW,
+        VAL_MAX_ROW, normalize, includeDateTime));
 
   const batchMeanAbsoluteErrors = [];
   const batchSizes = [];
   await dataset.forEach(dataItem => {
-    const features = dataItem[0];
-    const targets = dataItem[1];
+    const features = dataItem.xs;
+    const targets = dataItem.ys;
     const timeSteps = features.shape[1];
     batchSizes.push(features.shape[0]);
     batchMeanAbsoluteErrors.push(tf.tidy(
         () => tf.losses.absoluteDifference(
             targets,
             features.gather([timeSteps - 1], 1).gather([1], 2).squeeze([2]))));
   });
+
   const meanAbsoluteError = tf.tidy(() => {
     const batchSizesTensor = tf.tensor1d(batchSizes);
     const batchMeanAbsoluteErrorsTensor = tf.stack(batchMeanAbsoluteErrors);
@@ -83,7 +85,7 @@ export async function getBaselineMeanAbsoluteError(
  * Build a linear-regression model for the temperature-prediction problem.
  *
  * @param {tf.Shape} inputShape Input shape (without the batch dimenson).
- * @returns {tf.Model} A TensorFlow.js tf.Model instance.
+ * @returns {tf.LayersModel} A TensorFlow.js tf.LayersModel instance.
  */
 function buildLinearRegressionModel(inputShape) {
   const model = tf.sequential();
@@ -102,9 +104,9 @@ function buildLinearRegressionModel(inputShape) {
  * @param {number} dropoutRate Dropout rate of an optional dropout layer
  *   inserted between the two dense layers of the MLP. Optional. If not
  *   specified, no dropout layers will be included in the MLP.
- * @returns {tf.Model} A TensorFlow.js tf.Model instance.
+ * @returns {tf.LayersModel} A TensorFlow.js tf.LayersModel instance.
  */
-function buildMLPModel(inputShape, kernelRegularizer, dropoutRate) {
+export function buildMLPModel(inputShape, kernelRegularizer, dropoutRate) {
   const model = tf.sequential();
   model.add(tf.layers.flatten({inputShape}));
   model.add(
@@ -120,9 +122,10 @@ function buildMLPModel(inputShape, kernelRegularizer, dropoutRate) {
  * Build a simpleRNN-based model for the temperature-prediction problem.
  *
  * @param {tf.Shape} inputShape Input shape (without the batch dimenson).
- * @returns {tf.Model} A TensorFlow.js model consisting of a simpleRNN layer.
+ * @returns {tf.LayersModel} A TensorFlow.js model consisting of a simpleRNN
+ *   layer.
  */
-function buildSimpleRNNModel(inputShape) {
+export function buildSimpleRNNModel(inputShape) {
   const model = tf.sequential();
   const rnnUnits = 32;
   model.add(tf.layers.simpleRNN({
@@ -139,9 +142,9 @@ function buildSimpleRNNModel(inputShape) {
  * @param {tf.Shape} inputShape Input shape (without the batch dimenson).
  * @param {number} dropout Optional input dropout rate
  * @param {number} recurrentDropout Optional recurrent dropout rate.
- * @returns {tf.Model} A TensorFlow.js GRU model.
+ * @returns {tf.LayersModel} A TensorFlow.js GRU model.
  */
-function buildGRUModel(inputShape, dropout, recurrentDropout) {
+export function buildGRUModel(inputShape, dropout, recurrentDropout) {
   // TODO(cais): Recurrent dropout is currently not fully working.
   //   Make it work and add a flag to train-rnn.js.
   const model = tf.sequential();
@@ -163,7 +166,7 @@ function buildGRUModel(inputShape, dropout, recurrentDropout) {
  * @param {number} numTimeSteps Number of time steps in each input.
  *   exapmle
  * @param {number} numFeatures Number of features (for each time step).
- * @returns A compiled instance of `tf.Model`.
+ * @returns A compiled instance of `tf.LayersModel`.
  */
 export function buildModel(modelType, numTimeSteps, numFeatures) {
   const inputShape = [numTimeSteps, numFeatures];
@@ -197,9 +200,9 @@ export function buildModel(modelType, numTimeSteps, numFeatures) {
 /**
  * Train a model on the Jena weather data.
  *
- * @param {tf.Model} model A compiled tf.Model object. It is expected to
- *   have a 3D input shape `[numExamples, timeSteps, numFeatures].` and an
- *   output shape `[numExamples, 1]` for predicting the temperature value.
+ * @param {tf.LayersModel} model A compiled tf.LayersModel object. It is
+ *   expected to have a 3D input shape `[numExamples, timeSteps, numFeatures].`
+ *   and an output shape `[numExamples, 1]` for predicting the temperature value.
  * @param {JenaWeatherData} jenaWeatherData A JenaWeatherData object.
  * @param {boolean} normalize Whether to used normalized data for training.
  * @param {boolean} includeDateTime Whether to include date and time features

jena-weather/models_test.js

@@ -0,0 +1,84 @@
+/*
+ * @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.
+ * =============================================================================
+ */
+
+import * as tf from '@tensorflow/tfjs-node';
+
+import {JenaWeatherData} from "./data";
+import {buildGRUModel, buildMLPModel, buildSimpleRNNModel, getBaselineMeanAbsoluteError} from "./models";
+
+describe('Model creation', () => {
+  it('MLP', () => {
+    const model = buildMLPModel([8, 9]);
+    const config = model.layers[1].getConfig();
+    expect(config.kernelRegularizer).toEqual(null);
+    expect(model.inputs.length).toEqual(1);
+    expect(model.inputs[0].shape).toEqual([null, 8, 9]);
+    expect(model.outputs.length).toEqual(1);
+    expect(model.outputs[0].shape).toEqual([null, 1]);
+  });
+
+  it('MLP with kernel regularizer', () => {
+    const model = buildMLPModel([8, 9], tf.regularizers.l2({l2: 5e-2}));
+    const config = model.layers[1].getConfig();
+    expect(config.kernelRegularizer.config.l2).toEqual(5e-2);
+    expect(model.inputs.length).toEqual(1);
+    expect(model.inputs[0].shape).toEqual([null, 8, 9]);
+    expect(model.outputs.length).toEqual(1);
+    expect(model.outputs[0].shape).toEqual([null, 1]);
+  });
+
+  it('MLP with dropout', () => {
+    const model = buildMLPModel([8, 9], null, 0.5);
+    const denseConfig = model.layers[1].getConfig();
+    expect(denseConfig.kernelRegularize).toEqual(undefined);
+    const dropoutConfig = model.layers[model.layers.length - 2].getConfig();
+    expect(dropoutConfig.rate).toEqual(0.5);
+    expect(model.inputs.length).toEqual(1);
+    expect(model.inputs[0].shape).toEqual([null, 8, 9]);
+    expect(model.outputs.length).toEqual(1);
+    expect(model.outputs[0].shape).toEqual([null, 1]);
+  });
+});
+
+describe('RNN', () => {
+  it('simpleRNN', () => {
+    const model = buildSimpleRNNModel([8, 9]);
+    expect(model.inputs.length).toEqual(1);
+    expect(model.inputs[0].shape).toEqual([null, 8, 9]);
+    expect(model.outputs.length).toEqual(1);
+    expect(model.outputs[0].shape).toEqual([null, 1]);
+  });
+
+  it('buildGRUModel', () => {
+    const model = buildGRUModel([8, 9]);
+    expect(model.inputs.length).toEqual(1);
+    expect(model.inputs[0].shape).toEqual([null, 8, 9]);
+    expect(model.outputs.length).toEqual(1);
+    expect(model.outputs[0].shape).toEqual([null, 1]);
+  });
+});
+
+describe('getBaselineMeanAbsoluteError', () => {
+  it('getBaselineMeanAbsoluteError', async () => {
+    const dataset = new JenaWeatherData();
+    await dataset.load();
+
+    const baselineMAE = await getBaselineMeanAbsoluteError(
+        dataset, true, false, 10 * 24 * 6, 6, 24 * 6);
+    expect(baselineMAE).toBeCloseTo(0.29033);
+  });
+});

jena-weather/package.json

@@ -9,18 +9,19 @@
     "node": ">=8.9.0"
   },
   "dependencies": {
-    "@tensorflow/tfjs": "^1.0.2",
+    "@tensorflow/tfjs": "^1.1.0",
     "@tensorflow/tfjs-vis": "^1.0.3"
   },
   "scripts": {
     "watch": "cross-env NODE_ENV=development parcel index.html --no-hmr --open",
     "build": "cross-env NODE_ENV=production parcel build index.html --no-minify --public-url ./",
     "link-local": "yalc link",
+    "test": "babel-node run_tests.js",
     "train-rnn": "babel-node --max_old_space_size=4096 train-rnn.js"
   },
   "devDependencies": {
-    "@tensorflow/tfjs-node": "^1.0.2",
-    "@tensorflow/tfjs-node-gpu": "^1.0.2",
+    "@tensorflow/tfjs-node": "^1.1.0",
+    "@tensorflow/tfjs-node-gpu": "^1.1.0",
     "argparse": "^1.0.10",
     "babel-cli": "^6.26.0",
     "babel-core": "^6.26.3",

jena-weather/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']);

jena-weather/train-rnn.js

@@ -78,6 +78,11 @@ function parseArguments() {
   parser.addArgument(
       '--epochs',
       {type: 'int', defaultValue: 20, help: 'Number of training epochs'});
+  parser.addArgument( '--earlyStoppingPatience', {
+    type: 'int',
+    defaultValue: 2,
+    help: 'Optional patience number for EarlyStoppingCallback'
+   });
   parser.addArgument('--logDir', {
     type: 'string',
     help: 'Optional tensorboard log directory, to which the loss and ' +
@@ -121,14 +126,23 @@ async function main() {
     const model = buildModel(
         args.modelType, Math.floor(args.lookBack / args.step), numFeatures);
 
-    let callback = null;
+    let callback = [];
     if (args.logDir != null) {
       console.log(
           `Logging to tensorboard. ` +
           `Use the command below to bring up tensorboard server:\n` +
           `  tensorboard --logdir ${args.logDir}`);
-      callback = tfn.node.tensorBoard(
-          args.logDir, {updateFreq: args.logUpdateFreq});
+      callback.push(tfn.node.tensorBoard(args.logDir, {
+        updateFreq: args.logUpdateFreq
+      }));
+    }
+    if (args.earlyStoppingPatience != null) {
+      console.log(
+          `Using earlyStoppingCallback with patience ` +
+          `${args.earlyStoppingPatience}.`);
+      callback.push(tfn.callbacks.earlyStopping({
+        patience: args.earlyStoppingPatience
+      }));
     }
 
     await trainModel(
@@ -138,4 +152,6 @@ async function main() {
   }
 }
 
-main();
+if (require.main === module) {
+  main();
+}