Fix demos and models

README.md

@@ -25,9 +25,9 @@ See the [TypeScript starter project](https://github.com/PAIR-code/deeplearnjs/tr
 short example that sums an array with a scalar (broadcasted):
 
 ```ts
-import {Array1D, NDArrayMathGPU, Scalar} from 'deeplearn';
+import {Array1D, ENV, Scalar} from 'deeplearn';
 
-const math = new NDArrayMathGPU();
+const math = ENV.math;
 const a = Array1D.new([1, 2, 3]);
 const b = Scalar.new(2);
 
@@ -63,7 +63,7 @@ After importing the library, the API will be available as `dl` in the global
 namespace.
 
 ```js
-var math = new dl.NDArrayMathGPU();
+var math = dl.ENV.math;
 var a = dl.Array1D.new([1, 2, 3]);
 var b = dl.Scalar.new(2);
 

demos/adder/adder.ts

@@ -15,13 +15,13 @@
  * =============================================================================
  */
 
-import {NDArrayMathGPU, Scalar} from 'deeplearn';
+import {ENV, Scalar} from 'deeplearn';
 
 const outputElement = document.getElementById('output');
 const inA: HTMLInputElement = document.getElementById('A') as HTMLInputElement;
 const inB: HTMLInputElement = document.getElementById('B') as HTMLInputElement;
 
-const math = new NDArrayMathGPU();
+const math = ENV.math;
 
 export async function execute(event?: Event) {
   const a = Scalar.new(+inA.value);

demos/benchmarks/batchnormalization3d_benchmark.ts

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 // tslint:disable-next-line:max-line-length
-import {Array1D, Array3D, NDArrayMathCPU, NDArrayMathGPU} from 'deeplearn';
+import {Array1D, Array3D, ENV, NDArrayMath} from 'deeplearn';
 
 import {BenchmarkTest, LAST_RUN_CPU_CUTOFF_MS} from './benchmark';
 import * as benchmark_util from './benchmark_util';
@@ -29,7 +29,9 @@ export class BatchNormalization3DCPUBenchmark implements BenchmarkTest {
         resolve(-1);
       });
     }
-    const math = new NDArrayMathCPU();
+    const safeMode = false;
+    const math = new NDArrayMath('cpu', safeMode);
+    ENV.setMath(math);
     const x = Array3D.randUniform([size, size, 8], -1, 1);
     const mean = Array1D.new([0]);
     const variance = Array1D.new([1]);
@@ -48,7 +50,9 @@ export class BatchNormalization3DCPUBenchmark implements BenchmarkTest {
 
 export class BatchNormalization3DGPUBenchmark implements BenchmarkTest {
   async run(size: number) {
-    const math = new NDArrayMathGPU();
+    const safeMode = false;
+    const math = new NDArrayMath('webgl', safeMode);
+    ENV.setMath(math);
     const x = Array3D.randUniform([size, size, 8], -1, 1);
     const mean = Array1D.new([0]);
     const variance = Array1D.new([1]);

demos/benchmarks/benchmark_util.ts

@@ -1,35 +1,12 @@
 
-import {ENV, NDArray, NDArrayMathGPU} from 'deeplearn';
+import {NDArray, NDArrayMath} from 'deeplearn';
 
 export async function warmupAndBenchmarkGPU(
-    math: NDArrayMathGPU, benchmark: () => NDArray): Promise<number> {
-  const gpgpu = math.getGPGPUContext();
-
-  let out: NDArray;
-  const saveOutputBenchmark = () => {
-    out = benchmark();
-  };
-
+    math: NDArrayMath, benchmark: () => NDArray): Promise<number> {
   // Warmup.
-  if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_ENABLED')) {
-    await gpgpu.runQuery(saveOutputBenchmark);
-  } else {
-    saveOutputBenchmark();
-    out.dataSync();
-  }
-  out.dispose();
-
-  let totalTime: number;
-  if (ENV.get('WEBGL_DISJOINT_QUERY_TIMER_EXTENSION_RELIABLE')) {
-    totalTime = await gpgpu.runQuery(saveOutputBenchmark);
-  } else {
-    const start = performance.now();
-
-    saveOutputBenchmark();
-    out.dataSync();
-
-    totalTime = performance.now() - start;
-  }
+  const out = benchmark();
+  await out.data();
   out.dispose();
-  return totalTime;
+  // Real timing.
+  return math.time(benchmark);
 }

demos/benchmarks/benchmarks.ts

@@ -179,7 +179,7 @@ export class MathBenchmark extends MathBenchmarkPolymer {
     return runNumberRowElement;
   }
 
-  private runBenchmarkSteps(
+  private async runBenchmarkSteps(
       chart: Chart, runGroup: BenchmarkRunGroup, benchmarkRunGroupIndex: number,
       step: number) {
     const runNumbersTable =
@@ -207,7 +207,6 @@ export class MathBenchmark extends MathBenchmarkPolymer {
     const runNumberRowElement = document.createElement('div');
     runNumberRowElement.className = 'run-numbers-row math-benchmark';
 
-    const runPromises: Array<Promise<number>> = [];
     const rowValues: string[] = [step.toString()];
     for (let i = 0; i < runGroup.benchmarkRuns.length; i++) {
       const run = runGroup.benchmarkRuns[i];
@@ -218,48 +217,23 @@ export class MathBenchmark extends MathBenchmarkPolymer {
           step;
 
       const opType = runGroup.selectedOption;
-      runPromises.push(test.run(size, opType, runGroup.params[opType]));
-    }
-
-    Promise.all(runPromises).then(results => {
-      for (let i = 0; i < runGroup.benchmarkRuns.length; i++) {
-        const benchmarkRun = runGroup.benchmarkRuns[i];
-
-        const size = runGroup.stepToSizeTransformation != null ?
-            runGroup.stepToSizeTransformation(step) :
-            step;
-
-        let resultString: string;
-        let logString: string;
-        let time = 0;
-        let success = true;
-        try {
-          time = results[i];
-          resultString = time.toFixed(3) + 'ms';
-          logString = resultString;
-        } catch (e) {
-          success = false;
-          resultString = 'Error';
-          logString = e.message;
-        }
+      const time = await test.run(size, opType, runGroup.params[opType]);
+      const resultString = time.toFixed(3) + 'ms';
 
-        if (time >= 0) {
-          if (success) {
-            benchmarkRun.chartData.push({x: step, y: time});
-          }
-          rowValues.push(resultString);
-        }
-        console.log(`${benchmarkRun.name}[${size}]: ${logString}`);
+      if (time >= 0) {
+        run.chartData.push({x: step, y: time});
+        rowValues.push(resultString);
       }
-      runNumbersTable.appendChild(this.buildRunNumbersRow(rowValues));
-
-      step += runGroup.stepSize;
-      // Allow the UI to update.
-      setTimeout(
-          () => this.runBenchmarkSteps(
-              chart, runGroup, benchmarkRunGroupIndex, step),
-          100);
-    });
+      console.log(`${run.name}[${size}]: ${resultString}`);
+    }
+    runNumbersTable.appendChild(this.buildRunNumbersRow(rowValues));
+
+    step += runGroup.stepSize;
+    // Allow the UI to update.
+    setTimeout(
+        () => this.runBenchmarkSteps(
+            chart, runGroup, benchmarkRunGroupIndex, step),
+        100);
   }
 }
 document.registerElement(MathBenchmark.prototype.is, MathBenchmark);

demos/benchmarks/conv_benchmarks.ts

@@ -16,7 +16,7 @@
  */
 
 // tslint:disable-next-line:max-line-length
-import {Array1D, Array3D, Array4D, conv_util, NDArray, NDArrayMathGPU} from 'deeplearn';
+import {Array1D, Array3D, Array4D, conv_util, ENV, NDArray} from 'deeplearn';
 
 import {BenchmarkTest} from './benchmark';
 import * as benchmark_util from './benchmark_util';
@@ -34,7 +34,7 @@ export interface DepthwiseConvParams extends ConvParams { channelMul: number; }
 
 export class ConvGPUBenchmark implements BenchmarkTest {
   async run(size: number, opType: string, params: ConvParams): Promise<number> {
-    const math = new NDArrayMathGPU();
+    const math = ENV.math;
 
     const inDepth = params.inDepth;
     const inShape: [number, number, number] = [size, size, inDepth];

demos/benchmarks/matmul_benchmarks.ts

@@ -15,7 +15,8 @@
  * =============================================================================
  */
 // tslint:disable-next-line:max-line-length
-import {Array2D, NDArrayMathCPU, NDArrayMathGPU} from 'deeplearn';
+import {Array2D, ENV, NDArrayMath} from 'deeplearn';
+
 import {BenchmarkTest, LAST_RUN_CPU_CUTOFF_MS} from './benchmark';
 import * as benchmark_util from './benchmark_util';
 
@@ -27,7 +28,9 @@ export class MatmulCPUBenchmark implements BenchmarkTest {
         resolve(-1);
       });
     }
-    const math = new NDArrayMathCPU();
+    const safeMode = false;
+    const math = new NDArrayMath('cpu', safeMode);
+    ENV.setMath(math);
     const a = Array2D.randUniform([size, size], -1, 1);
     const b = Array2D.randUniform([size, size], -1, 1);
     const start = performance.now();
@@ -41,8 +44,9 @@ export class MatmulCPUBenchmark implements BenchmarkTest {
 
 export class MatmulGPUBenchmark implements BenchmarkTest {
   async run(size: number): Promise<number> {
-    const math = new NDArrayMathGPU();
-
+    const safeMode = false;
+    const math = new NDArrayMath('webgl', safeMode);
+    ENV.setMath(math);
     const a = Array2D.randNormal([size, size]);
     const b = Array2D.randNormal([size, size]);
 

demos/benchmarks/pool_benchmarks.ts

@@ -16,7 +16,7 @@
  */
 
 // tslint:disable-next-line:max-line-length
-import {Array3D, conv_util, NDArrayMath, NDArrayMathCPU, NDArrayMathGPU} from 'deeplearn';
+import {Array3D, conv_util, ENV, NDArrayMath} from 'deeplearn';
 import {BenchmarkTest} from './benchmark';
 import * as benchmark_util from './benchmark_util';
 
@@ -58,7 +58,9 @@ function getPoolingOp(option: string, math: NDArrayMath): (
 export class PoolCPUBenchmark implements BenchmarkTest {
   run(size: number, option: string,
       params: PoolBenchmarkParams): Promise<number> {
-    const math = new NDArrayMathCPU();
+    const safeMode = false;
+    const math = new NDArrayMath('cpu', safeMode);
+    ENV.setMath(math);
     const outputDepth = params.depth;
     const xShape: [number, number, number] = [size, size, outputDepth];
     const fieldSize = params.fieldSize;
@@ -83,8 +85,9 @@ export class PoolCPUBenchmark implements BenchmarkTest {
 export class PoolGPUBenchmark implements BenchmarkTest {
   async run(size: number, option: string, params: PoolBenchmarkParams):
       Promise<number> {
-    const math = new NDArrayMathGPU();
-
+    const safeMode = false;
+    const math = new NDArrayMath('webgl', safeMode);
+    ENV.setMath(math);
     const outputDepth = params.depth;
     const xShape: [number, number, number] = [size, size, outputDepth];
     const fieldSize = params.fieldSize;
@@ -93,11 +96,8 @@ export class PoolGPUBenchmark implements BenchmarkTest {
     const op = getPoolingOp(option, math);
 
     const benchmark = () => op(x, fieldSize, stride, 'same');
-
     const time = await benchmark_util.warmupAndBenchmarkGPU(math, benchmark);
-
     x.dispose();
-    math.dispose();
 
     return time;
   }

demos/benchmarks/reduction_ops_benchmark.ts

@@ -15,7 +15,7 @@
  * =============================================================================
  */
 // tslint:disable-next-line:max-line-length
-import {Array2D, NDArray, NDArrayMath, NDArrayMathCPU, NDArrayMathGPU, Scalar} from 'deeplearn';
+import {Array2D, ENV, NDArray, NDArrayMath, Scalar} from 'deeplearn';
 
 import {BenchmarkTest} from './benchmark';
 import * as benchmark_util from './benchmark_util';
@@ -42,7 +42,9 @@ function getReductionOp(option: string, math: NDArrayMath): (input: NDArray) =>
 
 export class ReductionOpsCPUBenchmark implements BenchmarkTest {
   async run(size: number, option: string): Promise<number> {
-    const math = new NDArrayMathCPU();
+    const safeMode = false;
+    const math = new NDArrayMath('cpu', safeMode);
+    ENV.setMath(math);
     const input = Array2D.randUniform([size, size], -1, 1);
     const op = getReductionOp(option, math);
     const start = performance.now();
@@ -58,7 +60,9 @@ export class ReductionOpsCPUBenchmark implements BenchmarkTest {
 
 export class ReductionOpsGPUBenchmark implements BenchmarkTest {
   async run(size: number, option: string) {
-    const math = new NDArrayMathGPU();
+    const safeMode = false;
+    const math = new NDArrayMath('webgl', safeMode);
+    ENV.setMath(math);
     const input = Array2D.randUniform([size, size], -1, 1);
     const op = getReductionOp(option, math);
 

demos/benchmarks/unary_ops_benchmark.ts

@@ -16,7 +16,8 @@
  */
 
 // tslint:disable-next-line:max-line-length
-import {Array2D, NDArray, NDArrayMath, NDArrayMathCPU, NDArrayMathGPU} from 'deeplearn';
+import {Array2D, ENV, NDArray, NDArrayMath} from 'deeplearn';
+
 import {BenchmarkTest} from './benchmark';
 import * as benchmark_util from './benchmark_util';
 
@@ -61,7 +62,9 @@ function getUnaryOp(option: string, math: NDArrayMath) {
 
 export class UnaryOpsCPUBenchmark implements BenchmarkTest {
   async run(size: number, option: string): Promise<number> {
-    const math = new NDArrayMathCPU();
+    const safeMode = false;
+    const math = new NDArrayMath('cpu', safeMode);
+    ENV.setMath(math);
     const input = Array2D.randUniform([size, size], -1, 1);
     const op = getUnaryOp(option, math);
     const start = performance.now();
@@ -77,7 +80,9 @@ export class UnaryOpsCPUBenchmark implements BenchmarkTest {
 
 export class UnaryOpsGPUBenchmark implements BenchmarkTest {
   async run(size: number, option: string) {
-    const math = new NDArrayMathGPU();
+    const safeMode = false;
+    const math = new NDArrayMath('webgl', safeMode);
+    ENV.setMath(math);
     const input = Array2D.randUniform([size, size], -1, 1);
     const op = getUnaryOp(option, math);