Add non-default noiseShape support to dropout op (#1782)

This PR is a followup PR for [tfjs-core#1343](#1343) which added `dropout` op, dropout feature initially requested in [tfjs#117](tensorflow/tfjs#117). This PR cleans up a TODO, `implements non default noise shape`. The implementation of noiseShape aligns with TensorFlow Python's dropout API.

The `noiseShape` feature would benefit further dropout related development, for example, make `tf.layers.dropout` support `noiseShape` configuration.

**Relative PR:**
* Make dropout layer support non default noiseShape and seed [tensorflow/tfjs-layers#556](tensorflow/tfjs-layers#556)

**Reference:**
* [TensorFlow tf.nn.dropout documentation](https://www.tensorflow.org/api_docs/python/tf/nn/dropout)
* [TensorFlow tf.nn.dropout implementation](https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/nn_ops.py#L2982-L3054)
* [TensorFlow _get_noise_shape implementation](https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/ops/nn_ops.py#L2903-L2925)

Author: syt123450

src/ops/dropout.ts

@@ -15,43 +15,61 @@
  * =============================================================================
  */
 
-import {Scalar, Tensor} from '../tensor';
-import {arraysEqual} from '../util';
+import {Tensor} from '../tensor';
+import {convertToTensor} from '../tensor_util_env';
+import {TensorLike} from '../types';
+import * as util from '../util';
 
 import {randomUniform} from './array_ops';
-import {sub} from './binary_ops';
+import {getNoiseShape} from './dropout_util';
 import {op} from './operation';
 
 /**
- * Sets entries in `x` to zero at random, while scaling the entire tensor.
+ * Computes dropout.
+ *
  * ```js
- * const x = tf.range(1, 21).reshape([10, 2]);
- * const rate = 0.5;
- * const seed = 23;
- * const noiseShape = null || x.shape;
- * const tensor = tf.dropout(x, rate, noiseShape, seed);
+ * const x = tf.tensor1d([1, 2, 2, 1]);
+ * const rate = 0.75;
+ * const output = tf.dropout(x, rate);
+ * output.print();
  * ```
- * @param x input tensor.
- * @param level fraction of the entries in the tensor that will be set to 0.
- * @param noiseShape shape of randomly generated keep/drop flags, must be
- *   broadcastable to the shape of `x`.
- * @param seed random seed to ensure determinism.
- * @returns Result of the dropout operation.
+ *
+ * @param x A floating point Tensor or TensorLike.
+ * @param rate A float in the range [0, 1). The probability that each element
+ *   of x is discarded.
+ * @param noiseShape An array of numbers of type int32, representing the
+ * shape for randomly generated keep/drop flags. If the noiseShape has null
+ * value, it will be automatically replaced with the x's relative dimension
+ * size. Optional.
+ * @param seed Used to create random seeds. Optional.
+ * @returns A Tensor of the same shape of x.
  */
+/** @doc {heading: 'Operations', subheading: 'Dropout'} */
 function dropout_(
-    x: Tensor, rate: Scalar|number, noiseShape?: number[],
-    seed?: number): Tensor {
-  if (noiseShape != null && !arraysEqual(x.shape, noiseShape)) {
-    // TODO(VariableVasasMT): implement non default noise shape
-    throw new Error(
-        'Non-default noise shape is not implemented yet: ' +
-        JSON.stringify(noiseShape));
+    x: Tensor|TensorLike, rate: number, noiseShape?: number[],
+    seed?: number|string): Tensor {
+  const $x = convertToTensor(x, 'x', 'dropout');
+
+  util.assert(
+      $x.dtype === 'float32',
+      () => `x has to be a floating point tensor since it's going to be ` +
+          `scaled, but got a ${$x.dtype} tensor instead.`);
+  util.assert(
+      rate >= 0 && rate < 1,
+      () => `rate must be a float in the range [0, 1), but got ${rate}.`);
+
+  if (rate === 0) {
+    return x instanceof Tensor ? $x.clone() : $x;
   }
 
-  let multiplier = randomUniform(x.shape, 0, 1, 'float32', seed).greater(rate);
-  // Scale the kept elements, so the expected sum is unchanged.
-  multiplier = multiplier.div(sub(1, rate) as Scalar);
-  return x.mul(multiplier);
+  const $noiseShape = getNoiseShape($x, noiseShape);
+  const keepProb = 1 - rate;
+  const multiplier = randomUniform($noiseShape, 0, 1, 'float32', seed)
+                         .add(keepProb)
+                         .floor()
+                         .div(keepProb);
+
+  return $x.mul(multiplier);
 }
 
 export const dropout = op({dropout_});

src/ops/dropout_test.ts

@@ -14,44 +14,161 @@
  * limitations under the License.
  * =============================================================================
  */
+
 import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
-import {Tensor} from '../tensor';
-
-function countParams(x: Tensor): number {
-  const shape = x.shape;
-  if (shape.length > 0) {
-    return shape.reduce((a: number, b: number) => a * b);
-  } else {
-    // Scalar.
-    return 1;
-  }
-}
+import {expectArraysClose} from '../test_util';
+
+import {tensor1d, tensor2d} from './tensor_ops';
 
 describeWithFlags('dropout', ALL_ENVS, () => {
-  const dropoutLevels = [0, 0.75];
-  const seed = 23;
-  for (const dropoutLevel of dropoutLevels) {
-    it(`Level = ${dropoutLevel}`, async () => {
-      const x = tf.range(1, 21).reshape([10, 2]);
-      const y = tf.dropout(x, tf.scalar(dropoutLevel), null, seed);
-      expect(y.dtype).toEqual(x.dtype);
-      expect(y.shape).toEqual(x.shape);
-      const xValue = await x.data();
-      const yValue = await y.data();
-      let nKept = 0;
-      for (let i = 0; i < xValue.length; ++i) {
-        if (yValue[i] !== 0) {
-          nKept++;
-          expect(yValue[i]).toBeCloseTo(1 / (1 - dropoutLevel) * xValue[i]);
+  it('x 1d array, rate 0', async () => {
+    const x = tensor1d([1, 2, 2, 1]);
+    const rate = 0;
+    const output = tf.dropout(x, rate);
+    expect(output.dtype).toEqual(x.dtype);
+    expect(output.shape).toEqual(x.shape);
+    expectArraysClose(await x.data(), await output.data());
+  });
+
+  it('x 1d array, rate 0.75', async () => {
+    const x = tensor1d([1, 2, 2, 1]);
+    const rate = 0.75;
+    const output = tf.dropout(x, rate);
+    expect(output.dtype).toEqual(x.dtype);
+    expect(output.shape).toEqual(x.shape);
+    const xValues = await x.data();
+    const outputValues = await output.data();
+    for (let i = 0; i < xValues.length; i++) {
+      if (outputValues[i] !== 0) {
+        expect(outputValues[i]).toBeCloseTo(1 / (1 - rate) * xValues[i]);
+      }
+    }
+  });
+
+  it('x 2d array, rate 0', async () => {
+    const x = tensor2d([1, 5, 2, 4, 3, 6], [2, 3]);
+    const rate = 0;
+    const output = tf.dropout(x, rate);
+    expect(output.dtype).toEqual(x.dtype);
+    expect(output.shape).toEqual(x.shape);
+    expectArraysClose(await x.data(), await output.data());
+  });
+
+  it('x 2d array, rate 0.75', async () => {
+    const x = tensor2d([1, 5, 2, 4, 3, 6], [2, 3]);
+    const rate = 0.75;
+    const output = tf.dropout(x, rate);
+    expect(output.dtype).toEqual(x.dtype);
+    expect(output.shape).toEqual(x.shape);
+    const xValues = await x.data();
+    const outputValues = await output.data();
+    for (let i = 0; i < xValues.length; i++) {
+      if (outputValues[i] !== 0) {
+        expect(outputValues[i]).toBeCloseTo(1 / (1 - rate) * xValues[i]);
+      }
+    }
+  });
+
+  it('x 1d array, rate 0.75, with noise shape length = 1', async () => {
+    const x = tensor1d([1, 2, 2, 1]);
+    const rate = 0.75;
+    const noiseShape = [1];
+    const output = tf.dropout(x, rate, noiseShape);
+    expect(output.dtype).toEqual(x.dtype);
+    expect(output.shape).toEqual(x.shape);
+    const xValues = await x.data();
+    const outputValues = await output.data();
+    const maskedOutput = outputValues[0];
+    for (let i = 0; i < xValues.length; i++) {
+      if (maskedOutput === 0) {
+        expect(outputValues[i]).toBe(maskedOutput);
+      }
+      if (outputValues[i] !== 0) {
+        expect(outputValues[i]).toBeCloseTo(1 / (1 - rate) * xValues[i]);
+      }
+    }
+  });
+
+  it('x 2d array, rate 0.75, with noise shape length = 2', async () => {
+    const x = tensor2d([1, 5, 2, 4, 3, 6], [2, 3]);
+    const rate = 0.75;
+    const noiseShape = [2, 1];
+    const output = tf.dropout(x, rate, noiseShape);
+    expect(output.dtype).toEqual(x.dtype);
+    expect(output.shape).toEqual(x.shape);
+    const xValues = await x.data();
+    const outputValues = await output.data();
+    for (let i = 0; i < x.shape[0]; i++) {
+      const maskedOutput = outputValues[i * x.shape[1]];
+      if (maskedOutput !== 0) {
+        expect(maskedOutput)
+            .toBeCloseTo(1 / (1 - rate) * xValues[i * x.shape[1]]);
+      } else {
+        for (let j = 0; j < x.shape[1]; j++) {
+          expect(outputValues[i * x.shape[1] + j]).toBe(maskedOutput);
         }
       }
-      const numel = countParams(x);
-      if (dropoutLevel === 0) {
-        expect(nKept).toEqual(numel);
+    }
+  });
+
+  it('broadcast noise shape', async () => {
+    const x = tensor2d([1, 5, 2, 4, 3, 6], [2, 3]);
+    const rate = 0.75;
+    // broadcast noise shape, same output as using noiseShape [2, 1]
+    const noiseShape = [1];
+    const output = tf.dropout(x, rate, noiseShape);
+    expect(output.dtype).toEqual(x.dtype);
+    expect(output.shape).toEqual(x.shape);
+    const xValues = await x.data();
+    const outputValues = await output.data();
+    for (let i = 0; i < x.shape[0]; i++) {
+      const maskedOutput = outputValues[i * x.shape[1]];
+      if (maskedOutput !== 0) {
+        expect(maskedOutput)
+            .toBeCloseTo(1 / (1 - rate) * xValues[i * x.shape[1]]);
       } else {
-        expect(nKept).toBeLessThan(numel);
+        for (let j = 0; j < x.shape[1]; j++) {
+          expect(outputValues[i * x.shape[1] + j]).toBe(maskedOutput);
+        }
       }
-    });
-  }
+    }
+  });
+
+  it('x 1d array, rate 0.75, with seed', async () => {
+    const x = tensor1d([1, 2, 2, 1]);
+    const rate = 0.75;
+    const seed = 23;
+    const output = tf.dropout(x, rate, null, seed);
+    expect(output.dtype).toEqual(x.dtype);
+    expect(output.shape).toEqual(x.shape);
+    const xValues = await x.data();
+    const outputValues = await output.data();
+    for (let i = 0; i < xValues.length; i++) {
+      if (outputValues[i] !== 0) {
+        expect(outputValues[i]).toBeCloseTo(1 / (1 - rate) * xValues[i]);
+      }
+    }
+  });
+
+  it('x TensorLike object', async () => {
+    const x = [1.0, 2.0, 2.0, 1.0];
+    const rate = 0;
+    const output = tf.dropout(x, rate);
+    expect(output.dtype).toEqual('float32');
+    expect(output.shape).toEqual([4]);
+    expectArraysClose(await output.data(), x);
+  });
+
+  it('throws when x.dtype != float32', async () => {
+    const x = tensor1d([1, 2, 2, 1], 'int32');
+    const rate = 0.75;
+    expect(() => tf.dropout(x, rate)).toThrowError();
+  });
+
+  it('throws when rate is not in the range [0, 1)', async () => {
+    const x = tensor1d([1, 2, 2, 1]);
+    const rate = 1.5;
+    expect(() => tf.dropout(x, rate)).toThrowError();
+  });
 });

src/ops/dropout_util.ts

@@ -0,0 +1,49 @@
+/**
+ * @license
+ * Copyright 2019 Google Inc. 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 {Tensor} from '../tensor';
+import * as util from '../util';
+
+/**
+ * Normalize noise shape based on provided tensor and noise shape.
+ *
+ * @param x Tensor.
+ * @param noiseShape The shape for the randomly generated keep/drop flags, as
+ *   an array of numbers. Optional.
+ * @returns Normalized noise shape.
+ */
+export function getNoiseShape(x: Tensor, noiseShape?: number[]): number[] {
+  if (noiseShape == null) {
+    return x.shape.slice();
+  }
+  if (util.arraysEqual(x.shape, noiseShape)) {
+    return noiseShape;
+  }
+  if (x.shape.length === noiseShape.length) {
+    const newDimension: number[] = [];
+    for (let i = 0; i < x.shape.length; i++) {
+      if (noiseShape[i] == null && x.shape[i] != null) {
+        newDimension.push(x.shape[i]);
+      } else {
+        newDimension.push(noiseShape[i]);
+      }
+    }
+    return newDimension;
+  }
+
+  return noiseShape;
+}

src/ops/dropout_util_test.ts

@@ -0,0 +1,56 @@
+/**
+ * @license
+ * Copyright 2019 Google Inc. 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 '../index';
+import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
+import {getNoiseShape} from './dropout_util';
+
+describeWithFlags('getNoiseShape', ALL_ENVS, () => {
+  it('x.shape == noiseShape', async () => {
+    const x = tf.ones([2, 3]);
+    const noiseShape = [2, 3];
+    const shape = getNoiseShape(x, noiseShape);
+    expect(shape).toEqual([2, 3]);
+  });
+
+  it('x.shape and noiseShape have same length, different value', async () => {
+    const x = tf.ones([2, 3]);
+    const noiseShape = [2, 1];
+    const shape = getNoiseShape(x, noiseShape);
+    expect(shape).toEqual([2, 1]);
+  });
+
+  it('noiseShape has null value', async () => {
+    const x = tf.ones([2, 3]);
+    const noiseShape = [2, null];
+    const shape = getNoiseShape(x, noiseShape);
+    expect(shape).toEqual([2, 3]);
+  });
+
+  it('x.shape and noiseShape has different length', async () => {
+    const x = tf.ones([2, 3, 4]);
+    const noiseShape = [2, 3];
+    const shape = getNoiseShape(x, noiseShape);
+    expect(shape).toEqual([2, 3]);
+  });
+
+  it('noiseShape is null', async () => {
+    const x = tf.ones([2, 3]);
+    const shape = getNoiseShape(x, null);
+    expect(shape).toEqual([2, 3]);
+  });
+});

src/tests.ts

@@ -58,6 +58,7 @@ import './ops/conv3d_test';
 import './ops/conv_util_test';
 import './ops/diag_test';
 import './ops/dropout_test';
+import './ops/dropout_util_test';
 import './ops/fused_test';
 import './ops/gather_nd_test';
 import './ops/image_ops_test';