From dba81e01cd285129afa1630c9a63c668e3c27e30 Mon Sep 17 00:00:00 2001
From: Leon Chen <leon@md.ai>
Date: Sun, 9 Oct 2016 13:31:23 -0400
Subject: [PATCH] update Model/Layer gpu togglin, update demos

---
 demos/src/mnist-cnn.js                    |  2 +-
 demos/src/mnist-vae.js                    |  2 +-
 demos/src/resnet50.css                    |  2 --
 demos/src/resnet50.js                     |  2 +-
 src/Layer.js                              | 14 ++++++++++++++
 src/Model.js                              | 16 ++++++++++++++++
 src/layers/convolutional/Convolution2D.js | 13 +++++++------
 7 files changed, 40 insertions(+), 11 deletions(-)

diff --git a/demos/src/mnist-cnn.js b/demos/src/mnist-cnn.js
index 7179d4c..7320417 100644
--- a/demos/src/mnist-cnn.js
+++ b/demos/src/mnist-cnn.js
@@ -119,7 +119,7 @@ export const MnistCnn = Vue.extend({
   methods: {
 
     toggleGpu: function () {
-      this.model.gpu = !this.useGpu
+      this.model.toggleGpu(!this.useGpu)
     },
 
     clear: function (e) {
diff --git a/demos/src/mnist-vae.js b/demos/src/mnist-vae.js
index db49cbd..9c9e2e0 100644
--- a/demos/src/mnist-vae.js
+++ b/demos/src/mnist-vae.js
@@ -92,7 +92,7 @@ export const MnistVae = Vue.extend({
   methods: {
 
     toggleGpu: function () {
-      this.model.gpu = !this.useGpu
+      this.model.toggleGpu(!this.useGpu)
     },
 
     activateCrosshairs: function (e) {
diff --git a/demos/src/resnet50.css b/demos/src/resnet50.css
index 8feced4..bccb629 100644
--- a/demos/src/resnet50.css
+++ b/demos/src/resnet50.css
@@ -81,8 +81,6 @@
       flex-direction: column;
       align-items: flex-start;
       justify-content: center;
-      user-select: none;
-      cursor: default;
 
       .output-class {
         display: flex;
diff --git a/demos/src/resnet50.js b/demos/src/resnet50.js
index c81a170..3188e7c 100644
--- a/demos/src/resnet50.js
+++ b/demos/src/resnet50.js
@@ -119,7 +119,7 @@ export const ResNet50 = Vue.extend({
   methods: {
 
     toggleGpu: function () {
-      this.model.gpu = !this.useGpu
+      this.model.toggleGpu(!this.useGpu)
     },
 
     imageURLInputChanged: function (e) {
diff --git a/src/Layer.js b/src/Layer.js
index 96c828c..e1a9b32 100644
--- a/src/Layer.js
+++ b/src/Layer.js
@@ -38,6 +38,20 @@ export default class Layer {
     })
   }
 
+  /**
+   * Toggle GPU mode on/off
+   * weblas must be available
+   * @param {boolean} mode - on/off
+   */
+  toggleGpu (mode) {
+    const newMode = typeof mode === 'undefined' ? !this._useWeblas : mode
+    if (newMode && weblas) {
+      this._useWeblas = true
+    } else {
+      this._useWeblas = false
+    }
+  }
+
   /**
    * Method for layer computational logic
    * @param {Tensor} x
diff --git a/src/Model.js b/src/Model.js
index 2b73e26..bba7efd 100644
--- a/src/Model.js
+++ b/src/Model.js
@@ -287,6 +287,22 @@ export default class Model {
     })
   }
 
+  /**
+   * Toggle GPU mode on/off
+   * Iterate through all layers and set `gpu` attribute
+   * @param {boolean} mode - on/off
+   */
+  toggleGpu (mode) {
+    if (typeof mode === 'undefined') {
+      this.gpu = !this.gpu
+    } else {
+      this.gpu = mode
+    }
+    for (let layer of this.modelLayersMap.values()) {
+      layer.toggleGpu(this.gpu)
+    }
+  }
+
   /**
    * Async function for recursively traversing the DAG
    * Graph object is stored in `this.modelDAG`, keyed by layer name.
diff --git a/src/layers/convolutional/Convolution2D.js b/src/layers/convolutional/Convolution2D.js
index ce319d8..a704e52 100644
--- a/src/layers/convolutional/Convolution2D.js
+++ b/src/layers/convolutional/Convolution2D.js
@@ -204,6 +204,9 @@ export default class Convolution2D extends Layer {
       x.tensor = x.tensor.transpose(1, 2, 0)
     }
 
+    let logging = false
+    if (['conv1', 'res2c_branch2a', 'res4e_branch2b', 'res3d_branch2b', 'res2b_branch2a', 'res4d_branch2c', 'res3c_branch2c'].includes(this.name)) logging = true
+
     this._calcOutputShape(x)
     this._padInput(x)
 
@@ -213,7 +216,7 @@ export default class Convolution2D extends Layer {
       imColsMat.createWeblasTensor()
     }
     let endTime = performance.now()
-    if (this.name === 'conv1') console.log('imColsMat', endTime - startTime)
+    if (logging) console.log(this.name, 'imColsMat', endTime - startTime)
 
     startTime = performance.now()
     const nbFilter = this.kernelShape[0]
@@ -239,7 +242,7 @@ export default class Convolution2D extends Layer {
       gemm(matMul.tensor, imColsMat.tensor, this._wRowsMat.tensor, 1, 1)
     }
     endTime = performance.now()
-    if (this.name === 'conv1') console.log('gemm', endTime - startTime)
+    if (logging) console.log(this.name, 'gemm', endTime - startTime)
 
     startTime = performance.now()
     let output = new Tensor([], this.outputShape)
@@ -252,12 +255,10 @@ export default class Convolution2D extends Layer {
     }
     x.tensor = output.tensor
     endTime = performance.now()
-    if (this.name === 'conv1') console.log('createOutput', endTime - startTime)
+    if (logging) console.log(this.name, 'nbFilter', nbFilter)
+    if (logging) console.log(this.name, 'createOutput', endTime - startTime)
 
-    startTime = performance.now()
     this.activation(x)
-    endTime = performance.now()
-    if (this.name === 'conv1') console.log('activation', endTime - startTime)
 
     // convert back to th ordering if necessary
     if (this.dimOrdering === 'th') {