Improvements thanks to #1 feedback

shiffman · shiffman · commit 01e2dd54bb62 · 2018-05-11T10:51:23.000-04:00
Updating this in preparation of YouTube tutorial. Changes thanks to feedback from @nsthorat. I think if we are to create generic `NeuralNetwork` or `Classifier` classes in ml5, a lot of the work will inevitably involve data helper classes so that the end user can work with vanilla arrays and tensors are all created and managed internally by ml5.
diff --git a/01_XOR/neuralnetwork.js b/01_XOR/neuralnetwork.js
@@ -1,5 +1,21 @@
+// Data helper classes
+
+// class ml5.Data
+
+// ml5.Data?
+class Data {
+  // Need to deal with shape
+  constructor(data) {
+    this.xs = tf.tensor2d(data.inputs);
+    this.ys = tf.tensor2d(data.targets);
+  }
+}
+
+
+// Helper class for a Batch of Data: ml5.Batch
 class Batch {
   constructor() {
+    // Need to deal with shape
     // this.shape = ??;
     this.data = [];
   }
@@ -12,7 +28,7 @@ class Batch {
 
 class NeuralNetwork {
 
-  constructor(inputs, hidden, outputs) {
+  constructor(inputs, hidden, outputs, lr) {
     this.model = tf.sequential();
     const hiddenLayer = tf.layers.dense({
       units: hidden,
@@ -21,13 +37,14 @@ class NeuralNetwork {
     });
     const outputLayer = tf.layers.dense({
       units: outputs,
-      inputShape: [hidden],
+      // inferred
+      // inputShape: [hidden],
       activation: 'sigmoid'
     });
     this.model.add(hiddenLayer);
     this.model.add(outputLayer);
 
-    const LEARNING_RATE = 0.5;
+    const LEARNING_RATE = lr || 0.5;
     const optimizer = tf.train.sgd(LEARNING_RATE);
 
     this.model.compile({
@@ -38,28 +55,46 @@ class NeuralNetwork {
   }
 
   predict(inputs) {
-    if (inputs instanceof Batch) {
-      return tf.tidy(() => {
-        const xs = tf.tensor2d(inputs.data);
-        return this.model.predict(xs).dataSync();
-      });
+    return tf.tidy(() => {
+      let data;
+      if (inputs instanceof Batch) {
+        data = inputs.data;
+      } else {
+        data = [inputs];
+      }
+      const xs = tf.tensor2d(data);
+      return this.model.predict(xs).dataSync();
+    });
+  }
+
+  setTrainingData(data) {
+    if (data instanceof Data) {
+      this.trainingData = data;
     } else {
-      return tf.tidy(() => {
-        const xs = tf.tensor2d([inputs]);
-        return this.model.predict(xs).dataSync();
-      });
+      this.trainingData = new Data(data);
     }
   }
 
-  async train(data, epochs, callback) {
-    const xs = tf.tensor2d(data.inputs);
-    const ys = tf.tensor2d(data.targets);
+  async train(callback, epochs, data) {
+    let xs, ys;
+    if (data) {
+      xs = tf.tensor2d(data.inputs);
+      ys = tf.tensor2d(data.targets);
+    } else if (this.trainingData) {
+      xs = this.trainingData.xs;
+      ys = this.trainingData.ys;
+    } else {
+      console.log("I have no data!");
+      return;
+    }
     await this.model.fit(xs, ys, {
-      epochs: epochs,
+      epochs: epochs || 1,
       shuffle: true
     });
-    xs.dispose();
-    ys.dispose();
+    if (data) {
+      xs.dispose();
+      ys.dispose();
+    }
     callback();
   }
-}
+}
diff --git a/01_XOR/sketch.js b/01_XOR/sketch.js
@@ -18,20 +18,21 @@ let counter = 0;
 let training = true;
 
 function train() {
-  nn.train(data, 10, finished);
+  nn.train(finished);
 }
 
 function finished() {
   counter++;
   statusP.html('training pass: ' + counter + '<br>framerate: ' + floor(frameRate()));
-  setTimeout(train, 10);
+  train();
 }
 
 let statusP;
 
 function setup() {
   createCanvas(400, 400);
   nn = new NeuralNetwork(2, 2, 1);
+  nn.setTrainingData(data);
   train();
   statusP = createP('0');
 }
@@ -66,4 +67,4 @@ function draw() {
     }
   }
   // }
-}
+}

Original file line number	Diff line number	Diff line change
`@@ -18,20 +18,21 @@ let counter = 0;`
`18`	`18`	`let training = true;`
`19`	`19`
`20`	`20`	`function train() {`
`21`		`- nn.train(data, 10, finished);`
	`21`	`+ nn.train(finished);`
`22`	`22`	`}`
`23`	`23`
`24`	`24`	`function finished() {`
`25`	`25`	`counter++;`
`26`	`26`	`statusP.html('training pass: ' + counter + '<br>framerate: ' + floor(frameRate()));`
`27`		`- setTimeout(train, 10);`
	`27`	`+ train();`
`28`	`28`	`}`
`29`	`29`
`30`	`30`	`let statusP;`
`31`	`31`
`32`	`32`	`function setup() {`
`33`	`33`	`createCanvas(400, 400);`
`34`	`34`	`nn = new NeuralNetwork(2, 2, 1);`
	`35`	`+ nn.setTrainingData(data);`
`35`	`36`	`train();`
`36`	`37`	`statusP = createP('0');`
`37`	`38`	`}`
`@@ -66,4 +67,4 @@ function draw() {`
`66`	`67`	`}`
`67`	`68`	`}`
`68`	`69`	`// }`
`69`		`-}`
	`70`	`+}`