Add non-default noiseShape and seed support to dropout layer (#556)

This PR makes dropout layer support non-default `noiseShape` and `seed`, adds relative tests for these feature in `dropout` layer and backend's `dropout` op.

The `noiseShape` feature in `dropout` layer would benefit further development, for example, `SpatialDropout2D` layer, feature requested in [#1453](#1453). Send this PR to ensure there won't be any broken change for further `noiseShape` related features.

FEATURE

Author: syt123450

src/backend/tfjs_backend_test.ts

@@ -742,6 +742,44 @@ describeMathCPUAndGPU('dropout', () => {
       }
     });
   }
+
+  it('Level=0.75, with noiseShape', () => {
+    const x = tensor2d(range(1, 21), [10, 2]);
+    const level = 0.75;
+    const noiseShape = [10, 1];
+    const y = K.dropout(x, level, noiseShape);
+    expect(y.dtype).toEqual(x.dtype);
+    expect(y.shape).toEqual(x.shape);
+    const xValue = x.dataSync();
+    const yValue = y.dataSync();
+    let nKept = 0;
+    for (let i = 0; i < x.shape[0]; i++) {
+      const maskedValue = yValue[i * x.shape[1]];
+      for (let j = 0; j < x.shape[1]; j++) {
+        const indice = i * x.shape[1] + j;
+        if (maskedValue !== 0) {
+          nKept++;
+          expect(yValue[indice]).toBeCloseTo(1 / (1 - level) * xValue[indice]);
+        } else {
+          expect(yValue[indice]).toEqual(0);
+        }
+      }
+    }
+    const numel = K.countParams(x);
+    expect(nKept).toBeLessThan(numel);
+  });
+
+  it('Level=0.75, with seed', () => {
+    const x = tensor2d(range(1, 21), [10, 2]);
+    const level = 0.75;
+    const seed = 23;
+    const y = K.dropout(x, level, null, seed);
+    expect(y.dtype).toEqual(x.dtype);
+    expect(y.shape).toEqual(x.shape);
+    const yValuesExpected =
+        [0, 0, 12, 16, 0, 0, 0, 0, 0, 0, 0, 48, 52, 0, 0, 0, 68, 0, 76, 0];
+    expectTensorsClose(y, tensor2d(yValuesExpected, [10, 2]));
+  });
 });
 
 describeMathCPUAndGPU('biasAdd', () => {

src/layers/core.ts

@@ -18,7 +18,7 @@ import {Activation as ActivationFn, getActivation, serializeActivation} from '..
 import * as K from '../backend/tfjs_backend';
 import {Constraint, ConstraintIdentifier, getConstraint, serializeConstraint} from '../constraints';
 import {DisposeResult, InputSpec, Layer, LayerArgs} from '../engine/topology';
-import {NotImplementedError, ValueError} from '../errors';
+import {ValueError} from '../errors';
 import {getInitializer, Initializer, InitializerIdentifier, serializeInitializer} from '../initializers';
 import {ActivationIdentifier} from '../keras_format/activation_config';
 import {Shape} from '../keras_format/common';
@@ -69,11 +69,6 @@ export class Dropout extends Layer {
     // So that the scalar doesn't get tidied up between executions.
     this.noiseShape = args.noiseShape;
     this.seed = args.seed;
-    if (this.seed != null) {
-      throw new NotImplementedError(
-          'Non-default seed is not implemented in Dropout layer yet: ' +
-          this.seed);
-    }
     this.supportsMasking = true;
   }
 
@@ -94,12 +89,6 @@ export class Dropout extends Layer {
     return tidy(() => {
       this.invokeCallHook(inputs, kwargs);
       const input = getExactlyOneTensor(inputs);
-      if (this.noiseShape != null &&
-          !util.arraysEqual(input.shape, this.noiseShape)) {
-        throw new NotImplementedError(
-            'Non-default noise shape is not implemented in Dropout ' +
-            'layer yet: ' + JSON.stringify(this.noiseShape));
-      }
       if (0 < this.rate && this.rate < 1) {
         const training =
             kwargs['training'] == null ? false : kwargs['training'];

src/layers/core_test.ts

@@ -52,8 +52,7 @@ describeMathCPUAndGPU('Dropout Layer', () => {
     const inputShape = [2, 3, 4];
     const trainingValues = [false, true];
     const dropoutRates = [0, 0.5];
-    const noiseShapes = [null, inputShape];
-    // TODO(cais): test non-default noiseShapes once they are supported.
+    const noiseShapes = [null, inputShape, [2, 3, 1]];
 
     for (const training of trainingValues) {
       for (const rate of dropoutRates) {
@@ -69,13 +68,37 @@ describeMathCPUAndGPU('Dropout Layer', () => {
             const xValue = x.dataSync();
             const yValue = y.dataSync();
             let nKept = 0;
-            for (let i = 0; i < xValue.length; ++i) {
-              if (yValue[i] !== 0) {
-                nKept++;
-                if (training) {
-                  expect(yValue[i]).toBeCloseTo(1 / (1 - rate));
-                } else {
-                  expect(yValue[i]).toBeCloseTo(1);
+            if (noiseShape === noiseShapes[2]) {  // customized noiseShape
+              for (let i = 0; i < x.shape[0]; ++i) {
+                for (let j = 0; j < x.shape[1]; ++j) {
+                  const maskedValue =
+                      yValue[i * x.shape[1] * x.shape[2] + j * x.shape[2]];
+                  for (let k = 0; k < x.shape[2]; ++k) {
+                    const indice =
+                        i * x.shape[1] * x.shape[2] + j * x.shape[2] + k;
+                    if (training) {
+                      if (maskedValue === 0) {
+                        expect(yValue[indice]).toEqual(0);
+                      } else {
+                        nKept++;
+                        expect(yValue[indice]).toBeCloseTo(1 / (1 - rate));
+                      }
+                    } else {
+                      nKept++;
+                      expect(yValue[indice]).toEqual(1);
+                    }
+                  }
+                }
+              }
+            } else {  // default noiseShape
+              for (let i = 0; i < xValue.length; ++i) {
+                if (yValue[i] !== 0) {
+                  nKept++;
+                  if (training) {
+                    expect(yValue[i]).toBeCloseTo(1 / (1 - rate));
+                  } else {
+                    expect(yValue[i]).toBeCloseTo(1);
+                  }
                 }
               }
             }
@@ -90,6 +113,20 @@ describeMathCPUAndGPU('Dropout Layer', () => {
       }
     }
   });
+
+  describe('tensor with seed get specific value', () => {
+    const training = true;
+    const rate = 0.5;
+    const noiseShape = [2, 3, 4];
+    const x = ones([2, 3, 4]);
+    const seed = 23;
+    const dropoutLayer = tfl.layers.dropout({rate, noiseShape, seed});
+    const y = dropoutLayer.apply(x, {training}) as Tensor;
+    const yValuesExpected = [
+      0, 2, 2, 2, 0, 0, 2, 2, 0, 0, 2, 2, 2, 0, 0, 0, 2, 0, 2, 0, 2, 2, 2, 0
+    ];
+    expectTensorsClose(y, tensor3d(yValuesExpected, [2, 3, 4]));
+  });
 });
 
 describeMathCPU('Dense Layer: Symbolic', () => {