Compute min and max of numerical columns over a Dataset. (#714)

* Compute min and max of numerical columns over a Dataset.

* Adds a few simple stream utilities, e.g. streamFromIncrementing, forEach, and resolveFully

* Don't test with  'WEBGL_FLOAT_TEXTURE_ENABLED': false; numerical bug

Author: davidsoergel

src/contrib/data/dataset.ts

@@ -19,6 +19,7 @@
 import * as seedrandom from 'seedrandom';
 
 import {BatchDataset} from './batch_dataset';
+import {computeDatasetStatistics, DatasetStatistics} from './statistics';
 import {DataStream} from './streams/data_stream';
 import {streamFromConcatenated} from './streams/data_stream';
 import {streamFromFunction} from './streams/data_stream';
@@ -46,6 +47,43 @@ export abstract class Dataset {
    */
   abstract async getStream(): Promise<DataStream<DatasetElement>>;
 
+  // TODO(soergel): Make Datasets report whether repeated getStream() calls
+  // produce the same result (e.g., reading from a file) or different results
+  // (e.g., from the webcam).  Currently we don't make this distinction but it
+  // could be important for the user to know.
+  // abstract isDeterministic(): boolean;
+
+  // TODO(soergel): memoize computeStatistics()
+
+  /**
+   * Gathers statistics from a Dataset (or optionally from a sample).
+   *
+   * This obtains a stream from the Dataset and, by default, does a full pass
+   * to gather the statistics.
+   *
+   * Statistics may be computed over a sample.  However: simply taking the first
+   * n items from the stream may produce a poor estimate if the stream is
+   * ordered in some way.
+   *
+   * A truly random shuffle of the stream would of course solve this
+   * problem, but many streams do not allow for this, instead providing only a
+   * sliding-window shuffle.  A partially-randomized sample could be obtained by
+   * shuffling over a window followed by taking the first n samples (where n is
+   * smaller than the shuffle window size).  However there is little point in
+   * using that approach here, because the cost is likely dominated by obtaining
+   * the data.  Thus, once we have filled our shuffle buffer, we may as well use
+   * all of that data instead of sampling from it.
+   *
+   * @param sampleSize The number of examples to take from the (possibly
+   *   shuffled) stream.
+   * @param shuffleWindowSize The size of the shuffle window to use, if any.
+   *   (Not recommended, as described above).
+   */
+  async computeStatistics(sampleSize?: number, shuffleWindowSize?: number):
+      Promise<DatasetStatistics> {
+    return computeDatasetStatistics(this, sampleSize, shuffleWindowSize);
+  }
+
   /**
    * Filters this dataset according to `predicate`.
    *
@@ -163,6 +201,8 @@ export abstract class Dataset {
     });
   }
 
+  // TODO(soergel): deep sharded shuffle, where supported
+
   /**
    * Randomly shuffles the elements of this dataset.
    *

src/contrib/data/statistics.ts

@@ -0,0 +1,119 @@
+/**
+ * @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 {Scalar, Tensor} from '../../tensor';
+
+import {Dataset} from './dataset';
+import {ElementArray} from './types';
+
+// TODO(soergel): Flesh out collected statistics.
+// For numeric columns we should provide mean, stddev, histogram, etc.
+// For string columns we should provide a vocabulary (at least, top-k), maybe a
+// length histogram, etc.
+// Collecting only numeric min and max is just the bare minimum for now.
+
+export type NumericColumnStatistics = {
+  min: number; max: number;
+};
+
+export type DatasetStatistics = {
+  [key: string]: NumericColumnStatistics
+};
+
+/**
+ * Provides a function that scales numeric values into the [0, 1] interval.
+ *
+ * @param min the lower bound of the inputs, which should be mapped to 0.
+ * @param max the upper bound of the inputs, which should be mapped to 1,
+ * @return A function that maps an input ElementArray to a scaled ElementArray.
+ */
+export function scaleTo01(min: number, max: number): (value: ElementArray) =>
+    ElementArray {
+  const range = max - min;
+  const minTensor: Tensor = Scalar.new(min);
+  const rangeTensor: Tensor = Scalar.new(range);
+  return (value: ElementArray): ElementArray => {
+    if (typeof (value) === 'string') {
+      throw new Error('Can\'t scale a string.');
+    } else {
+      if (value instanceof Tensor) {
+        const result = value.sub(minTensor).div(rangeTensor);
+        return result;
+      } else if (value instanceof Array) {
+        return value.map(v => (v - min) / range);
+      } else {
+        return (value - min) / range;
+      }
+    }
+  };
+}
+
+export async function computeDatasetStatistics(
+    dataset: Dataset, sampleSize?: number,
+    shuffleWindowSize?: number): Promise<DatasetStatistics> {
+  let stream = await dataset.getStream();
+  // TODO(soergel): allow for deep shuffle where possible.
+  if (shuffleWindowSize != null) {
+    stream = stream.shuffle(shuffleWindowSize);
+  }
+  if (sampleSize != null) {
+    stream = stream.take(sampleSize);
+  }
+
+  // TODO(soergel): prepare the column objects based on a schema.q
+  const result: DatasetStatistics = {};
+
+  await stream.forEach(e => {
+    for (const key in e) {
+      const value = e[key];
+      if (typeof (value) === 'string') {
+      } else {
+        let recordMin: number;
+        let recordMax: number;
+        if (value instanceof Tensor) {
+          recordMin = value.min().dataSync()[0];
+          recordMax = value.max().dataSync()[0];
+        } else if (value instanceof Array) {
+          recordMin = value.reduce((a, b) => Math.min(a, b));
+          recordMax = value.reduce((a, b) => Math.max(a, b));
+        } else if (!isNaN(value) && isFinite(value)) {
+          recordMin = value;
+          recordMax = value;
+        } else {
+          // TODO(soergel): don't throw; instead record the stats as "unknown".
+          throw new Error(`Cannot compute statistics: ${key} = ${value}`);
+        }
+        let columnStats: NumericColumnStatistics = result[key];
+        if (columnStats == null) {
+          columnStats = {
+            min: Number.POSITIVE_INFINITY,
+            max: Number.NEGATIVE_INFINITY
+          };
+          result[key] = columnStats;
+        }
+        columnStats.min = Math.min(columnStats.min, recordMin);
+        columnStats.max = Math.max(columnStats.max, recordMax);
+      }
+    }
+    // Returning undefined or null (i.e, type void) would indicate that the
+    // stream is exhausted.  So, we have to return *something* in order for
+    // resolveFully() to operate.
+    return {};
+  });
+  return result;
+}

src/contrib/data/statistics_test.ts

@@ -0,0 +1,55 @@
+/**
+ * @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 {ALL_FLOAT_ENVS, describeWithFlags} from '../../test_util';
+
+import {TestDataset} from './dataset_test';
+import {scaleTo01} from './statistics';
+
+describeWithFlags('makeDatasetStatistics', ALL_FLOAT_ENVS, () => {
+  it('computes numeric min and max over numbers, arrays, and Tensors', done => {
+    const ds = new TestDataset().skip(55);
+    ds.computeStatistics()
+        .then(stats => {
+          expect(stats['number'].min).toEqual(55);
+          expect(stats['number'].max).toEqual(99);
+          // The TestDataset includes cubes of the indices
+          expect(stats['numberArray'].min).toEqual(55);
+          expect(stats['numberArray'].max).toEqual(99 * 99 * 99);
+          expect(stats['Tensor'].min).toEqual(55);
+          expect(stats['Tensor'].max).toEqual(99 * 99 * 99);
+        })
+        .then(done)
+        .catch(done.fail);
+  });
+});
+
+describeWithFlags('scaleTo01', ALL_FLOAT_ENVS, () => {
+  it('scales numeric data to the [0, 1] interval', done => {
+    const ds = new TestDataset().skip(55);
+    const scaleFn = scaleTo01(55, 99 * 99 * 99);
+    const scaledDataset = ds.map(x => ({'Tensor': scaleFn(x['Tensor'])}));
+
+    scaledDataset.computeStatistics()
+        .then(stats => {
+          expect(stats['Tensor'].min).toBeCloseTo(0);
+          expect(stats['Tensor'].max).toBeCloseTo(1);
+        })
+        .then(done)
+        .catch(done.fail);
+  });
+});

src/contrib/data/streams/data_stream.ts

@@ -34,6 +34,14 @@ export function streamFromItems<T>(items: T[]): DataStream<T> {
   return new ArrayStream(items);
 }
 
+/**
+ * Create a `DataStream` of incrementing integers.
+ */
+export function streamFromIncrementing(start: number): DataStream<number> {
+  let i = start;
+  return streamFromFunction(() => i++);
+}
+
 /**
  * Create a `DataStream` from a function.
  */
@@ -95,13 +103,29 @@ export abstract class DataStream<T> {
   async collectRemaining(): Promise<T[]> {
     const result: T[] = [];
     let x = await this.next();
-    while (x !== undefined) {
+    while (x != null) {
       result.push(x);
       x = await this.next();
     }
     return result;
   }
 
+  /**
+   * Draw items from the stream until it is exhausted.
+   *
+   * This can be useful when the stream has side effects but no output.  In
+   * that case, calling this function guarantees that the stream will be fully
+   * processed.
+   */
+  async resolveFully(): Promise<void> {
+    let x = await this.next();
+    while (x != null) {
+      x = await this.next();
+    }
+  }
+
+  // TODO(soergel): Implement reduce() etc.
+
   /**
    * Filters this stream according to `predicate`.
    *
@@ -126,6 +150,15 @@ export abstract class DataStream<T> {
     return new MapStream(this, transform);
   }
 
+  /**
+   * Apply a function to every element of the stream.
+   *
+   * @param f A function to apply to each stream element.
+   */
+  async forEach(f: (value: T) => {}|Promise<{}>): Promise<void> {
+    return this.map(f).resolveFully();
+  }
+
   /**
    * Groups elements into batches.
    *
@@ -185,6 +218,8 @@ export abstract class DataStream<T> {
     return new PrefetchStream(this, bufferSize);
   }
 
+  // TODO(soergel): deep sharded shuffle, where supported
+
   /**
    * Randomly shuffles the elements of this stream.
    *
@@ -193,8 +228,8 @@ export abstract class DataStream<T> {
    * @param seed: (Optional.) An integer specifying the random seed that will
    *   be used to create the distribution.
    */
-  shuffle(bufferSize: number, seed?: string): DataStream<T> {
-    return new ShuffleStream(this, bufferSize, seed);
+  shuffle(windowSize: number, seed?: string): DataStream<T> {
+    return new ShuffleStream(this, windowSize, seed);
   }
 }
 
@@ -486,9 +521,9 @@ export class ShuffleStream<T> extends PrefetchStream<T> {
   private upstreamExhausted = false;
 
   constructor(
-      protected upstream: DataStream<T>, protected bufferSize: number,
+      protected upstream: DataStream<T>, protected windowSize: number,
       seed?: string) {
-    super(upstream, bufferSize);
+    super(upstream, windowSize);
     this.random = seedrandom(seed);
   }
 

src/contrib/data/streams/data_stream_test.ts

@@ -16,7 +16,7 @@
  * =============================================================================
  */
 
-import {DataStream} from './data_stream';
+import {DataStream, streamFromIncrementing} from './data_stream';
 import {streamFromConcatenated} from './data_stream';
 import {streamFromConcatenatedFunction} from './data_stream';
 import {streamFromFunction, streamFromItems} from './data_stream';
@@ -193,6 +193,16 @@ describe('DataStream', () => {
         .catch(done.fail);
   });
 
+  it('can be created with incrementing integers', done => {
+    const readStream = streamFromIncrementing(0).take(7);
+    readStream.collectRemaining()
+        .then(result => {
+          expect(result).toEqual([0, 1, 2, 3, 4, 5, 6]);
+        })
+        .then(done)
+        .catch(done.fail);
+  });
+
   it('can be concatenated', done => {
     const a = streamFromItems([1, 2, 3]);
     const b = streamFromItems([4, 5, 6]);

src/test_util.ts

@@ -22,16 +22,21 @@ import {Tensor} from './tensor';
 import {DataType, TypedArray} from './types';
 import * as util from './util';
 
-export const WEBGL_ENVS: Features[] = [
+const WEBGL_FLOAT_ENVS: Features[] = [
   {'BACKEND': 'webgl', 'WEBGL_FLOAT_TEXTURE_ENABLED': true, 'WEBGL_VERSION': 1},
-  {'BACKEND': 'webgl', 'WEBGL_FLOAT_TEXTURE_ENABLED': true, 'WEBGL_VERSION': 2},
   {
     'BACKEND': 'webgl',
-    'WEBGL_FLOAT_TEXTURE_ENABLED': false,
-    'WEBGL_VERSION': 1
-  },
+    'WEBGL_FLOAT_TEXTURE_ENABLED': true,
+    'WEBGL_VERSION': 2
+  }
 ];
+export const WEBGL_ENVS = WEBGL_FLOAT_ENVS.concat([{
+  'BACKEND': 'webgl',
+  'WEBGL_FLOAT_TEXTURE_ENABLED': false,
+  'WEBGL_VERSION': 1
+}]);
 export const CPU_ENVS: Features[] = [{'BACKEND': 'cpu'}];
+export const ALL_FLOAT_ENVS = WEBGL_FLOAT_ENVS.concat(CPU_ENVS);
 export const ALL_ENVS = WEBGL_ENVS.concat(CPU_ENVS);
 
 /** Accuracy for tests. */