Don't dispose input tensor when calling dl.variable(). (#712)

Nikhil Thorat · web-flow · commit 8326870e1d5c · 2018-02-13T22:48:18.000-05:00
diff --git a/src/optimizers/adadelta_optimizer.ts b/src/optimizers/adadelta_optimizer.ts
@@ -55,13 +55,17 @@ export class AdadeltaOptimizer extends Optimizer {
       const value = ENV.engine.registeredVariables[variableName];
       if (this.accumulatedGrads[variableName] == null) {
         const trainable = false;
-        this.accumulatedGrads[variableName] =
-            variable(zerosLike(value), trainable);
+        tidy(() => {
+          this.accumulatedGrads[variableName] =
+              variable(zerosLike(value), trainable);
+        });
       }
       if (this.accumulatedUpdates[variableName] == null) {
         const trainable = false;
-        this.accumulatedUpdates[variableName] =
-            variable(zerosLike(value), trainable);
+        tidy(() => {
+          this.accumulatedUpdates[variableName] =
+              variable(zerosLike(value), trainable);
+        });
       }
 
       const gradient = variableGradients[variableName];
diff --git a/src/optimizers/adadelta_optimizer_test.ts b/src/optimizers/adadelta_optimizer_test.ts
@@ -76,8 +76,8 @@ describeWithFlags('AdadeltaOptimizer', ALL_ENVS, () => {
     x.dispose();
     optimizer.dispose();
 
-    // There should be no more Tensors.
-    expect(dl.memory().numTensors).toBe(0);
+    // The only tensor remaining is the argument to variable().
+    expect(dl.memory().numTensors).toBe(1);
   });
 
   it('graph', () => {
diff --git a/src/optimizers/adagrad_optimizer.ts b/src/optimizers/adagrad_optimizer.ts
@@ -50,8 +50,10 @@ export class AdagradOptimizer extends Optimizer {
       const value = ENV.engine.registeredVariables[variableName];
       if (this.accumulatedGrads[variableName] == null) {
         const trainable = false;
-        this.accumulatedGrads[variableName] = variable(
-            fill(value.shape, this.initialAccumulatorValue), trainable);
+        tidy(() => {
+          this.accumulatedGrads[variableName] = variable(
+              fill(value.shape, this.initialAccumulatorValue), trainable);
+        });
       }
 
       const gradient = variableGradients[variableName];
diff --git a/src/optimizers/adagrad_optimizer_test.ts b/src/optimizers/adagrad_optimizer_test.ts
@@ -68,8 +68,8 @@ describeWithFlags('AdagradOptimizer', ALL_ENVS, () => {
     x.dispose();
     optimizer.dispose();
 
-    // There should be no more Tensors.
-    expect(dl.memory().numTensors).toBe(0);
+    // The only tensor remaining is the argument to variable().
+    expect(dl.memory().numTensors).toBe(1);
   });
 
   it('graph', () => {
diff --git a/src/optimizers/momentum_optimizer.ts b/src/optimizers/momentum_optimizer.ts
@@ -46,8 +46,10 @@ export class MomentumOptimizer extends SGDOptimizer {
       const value = ENV.engine.registeredVariables[variableName];
       if (this.accumulations[variableName] == null) {
         const trainable = false;
-        this.accumulations[variableName] =
-            variable(zerosLike(value), trainable);
+        tidy(() => {
+          this.accumulations[variableName] =
+              variable(zerosLike(value), trainable);
+        });
       }
 
       const accumulation = this.accumulations[variableName];
diff --git a/src/optimizers/momentum_optimizer_test.ts b/src/optimizers/momentum_optimizer_test.ts
@@ -67,8 +67,8 @@ describeWithFlags('MomentumOptimizer', ALL_ENVS, () => {
     x.dispose();
     optimizer.dispose();
 
-    // There should be no more Tensors.
-    expect(dl.memory().numTensors).toBe(0);
+    // The only tensor remaining is the argument to variable().
+    expect(dl.memory().numTensors).toBe(1);
   });
 
   it('graph', () => {
diff --git a/src/optimizers/optimizer_test.ts b/src/optimizers/optimizer_test.ts
@@ -68,8 +68,8 @@ describeWithFlags('optimizer', ALL_ENVS, () => {
     x.dispose();
     bias.dispose();
     strayVariable.dispose();
-    // There should be no more Tensors.
-    expect(dl.memory().numTensors).toBe(0);
+    // The only tensors remaining are the arguments to variable().
+    expect(dl.memory().numTensors).toBe(3);
   });
 
   it('varList array of all variables', () => {
diff --git a/src/optimizers/rmsprop_optimizer.ts b/src/optimizers/rmsprop_optimizer.ts
@@ -58,13 +58,17 @@ export class RMSPropOptimizer extends Optimizer {
       const value = ENV.engine.registeredVariables[variableName];
       if (this.accumulatedMeanSquares[variableName] == null) {
         const trainable = false;
-        this.accumulatedMeanSquares[variableName] =
-            variable(zerosLike(value), trainable);
+        tidy(() => {
+          this.accumulatedMeanSquares[variableName] =
+              variable(zerosLike(value), trainable);
+        });
       }
       if (this.accumulatedMoments[variableName] == null) {
         const trainable = false;
-        this.accumulatedMoments[variableName] =
-            variable(zerosLike(value), trainable);
+        tidy(() => {
+          this.accumulatedMoments[variableName] =
+              variable(zerosLike(value), trainable);
+        });
       }
 
       const accumulatedMeanSquare = this.accumulatedMeanSquares[variableName];
diff --git a/src/optimizers/rmsprop_optimizer_test.ts b/src/optimizers/rmsprop_optimizer_test.ts
@@ -78,7 +78,8 @@ describeWithFlags('RMSPropOptimizer', ALL_ENVS, () => {
 
     x.dispose();
     optimizer.dispose();
-    expect(dl.memory().numTensors).toBe(0);
+    // The only tensor remaining is the argument to variable().
+    expect(dl.memory().numTensors).toBe(1);
   });
 
   it('graph', () => {
diff --git a/src/optimizers/sgd_optimizer_test.ts b/src/optimizers/sgd_optimizer_test.ts
@@ -54,8 +54,8 @@ describeWithFlags('SGDOptimizer', ALL_ENVS, () => {
 
     optimizer.dispose();
     x.dispose();
-    // There should be no more Tensors.
-    expect(dl.memory().numTensors).toBe(0);
+    // The only tensor remaining is the argument to variable().
+    expect(dl.memory().numTensors).toBe(1);
   });
 
   it('graph', () => {
diff --git a/src/tensor.ts b/src/tensor.ts
@@ -945,7 +945,6 @@ export class Variable<R extends Rank = Rank> extends Tensor<R> {
     super(
         initialValue.shape, initialValue.dtype, null /* values */,
         initialValue.dataId);
-    initialValue.dispose();
     this.name = name;
     if (this.name == null) {
       this.name = Variable.nextVarId.toString();
diff --git a/src/variable_test.ts b/src/variable_test.ts
@@ -16,9 +16,9 @@
  */
 
 import * as dl from './index';
-import {ALL_ENVS, describeWithFlags, expectArraysClose} from './test_util';
 // tslint:disable-next-line:max-line-length
 import {Scalar, Tensor, Tensor1D, Tensor2D, Tensor3D, Tensor4D, variable, Variable} from './tensor';
+import {ALL_ENVS, describeWithFlags, expectArraysClose} from './test_util';
 import {Rank} from './types';
 
 describeWithFlags('variable', ALL_ENVS, () => {
@@ -66,7 +66,7 @@ describeWithFlags('variable', ALL_ENVS, () => {
       expect(dl.memory().numTensors).toBe(1);
 
       v = variable(value);
-      expect(dl.memory().numTensors).toBe(1);
+      expect(dl.memory().numTensors).toBe(2);
     });
 
     expect(dl.memory().numTensors).toBe(1);
@@ -76,6 +76,16 @@ describeWithFlags('variable', ALL_ENVS, () => {
     expect(dl.memory().numTensors).toBe(0);
   });
 
+  it('constructor does not dispose', () => {
+    const a = dl.scalar(2);
+    const v = dl.variable(a);
+
+    expect(dl.memory().numTensors).toBe(2);
+    expect(dl.memory().numDataBuffers).toBe(1);
+    expectArraysClose(v, [2]);
+    expectArraysClose(a, [2]);
+  });
+
   it('variables are assignable to tensors', () => {
     // This test asserts compilation, not doing any run-time assertion.
     const x0: Variable<Rank.R0> = null;
@@ -103,23 +113,30 @@ describeWithFlags('variable', ALL_ENVS, () => {
     expect(yh).toBeNull();
   });
 
-  it('assign will dispose old data', () => {
+  it('assign does not dispose old data', () => {
     let v: Variable<Rank.R1>;
     v = variable(dl.tensor1d([1, 2, 3]));
-    expect(dl.memory().numTensors).toBe(1);
+
+    expect(dl.memory().numTensors).toBe(2);
+    expect(dl.memory().numDataBuffers).toBe(1);
+
     expectArraysClose(v, [1, 2, 3]);
 
     const secondArray = dl.tensor1d([4, 5, 6]);
-    expect(dl.memory().numTensors).toBe(2);
+    expect(dl.memory().numTensors).toBe(3);
+    expect(dl.memory().numDataBuffers).toBe(2);
 
     v.assign(secondArray);
     expectArraysClose(v, [4, 5, 6]);
     // Assign doesn't dispose the 1st array.
-    expect(dl.memory().numTensors).toBe(2);
+    expect(dl.memory().numTensors).toBe(3);
+    expect(dl.memory().numDataBuffers).toBe(2);
 
     v.dispose();
-    // Disposing the variable disposes itself.
-    expect(dl.memory().numTensors).toBe(1);
+    // Disposing the variable disposes itself. The input to variable and
+    // secondArray are the only remaining tensors.
+    expect(dl.memory().numTensors).toBe(2);
+    expect(dl.memory().numDataBuffers).toBe(2);
   });
 
   it('shape must match', () => {
