Add serialization infrastructure to core (#997)

Moves the serialization infrastructure from -layers to -core. Places all the symbols under a second level export 'serialization'.
tensorflow · May 3, 2018 · bd82d33 · bd82d33
1 parent 316d5da
commit bd82d33
Show file tree

Hide file tree

Showing 17 changed files with 379 additions and 80 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -26,6 +26,7 @@ import * as environment from './environment';
 import {Environment} from './environment';
 // Serialization.
 import * as io from './io/io';
+import * as serialization from './serialization';
 import * as test_util from './test_util';
 import * as util from './util';
 import {version} from './version';
@@ -63,7 +64,7 @@ export {doc} from './doc';
 export const nextFrame = BrowserUtil.nextFrame;
 
 // Second level exports.
-export {environment, io, test_util, util, webgl};
+export {environment, io, serialization, test_util, util, webgl};
 
 // Backend specific.
 export {KernelBackend, BackendTimingInfo} from './kernels/backend';
diff --git a/src/optimizers/adadelta_optimizer.ts b/src/optimizers/adadelta_optimizer.ts
@@ -18,25 +18,31 @@
 import {ENV} from '../environment';
 import {keep, tidy} from '../globals';
 import {scalar, zerosLike} from '../ops/ops';
+// tslint:disable-next-line:max-line-length
+import {ConfigDict, Serializable, SerializableConstructor, SerializationMap} from '../serialization';
 import {Scalar} from '../tensor';
 import {NamedVariableMap} from '../types';
+
 import {Optimizer} from './optimizer';
 
 /** @doclink Optimizer */
 export class AdadeltaOptimizer extends Optimizer {
+  static className = 'AdadeltaOptimizer';
   private c: Scalar;
-  private epsilon: Scalar;
-  private rho: Scalar;
+  private epsilonScalar: Scalar;
+  private rhoScalar: Scalar;
   private oneMinusRho: Scalar;
 
   private accumulatedGrads: NamedVariableMap = {};
   private accumulatedUpdates: NamedVariableMap = {};
 
-  constructor(learningRate: number, rho: number, epsilon = 1e-8) {
+  constructor(
+      protected learningRate: number, protected rho: number,
+      protected epsilon = 1e-8) {
     super();
     this.c = keep(scalar(-learningRate));
-    this.epsilon = keep(scalar(epsilon));
-    this.rho = keep(scalar(rho));
+    this.epsilonScalar = keep(scalar(epsilon));
+    this.rhoScalar = keep(scalar(rho));
     this.oneMinusRho = keep(scalar(1 - rho));
   }
 
@@ -64,16 +70,16 @@ export class AdadeltaOptimizer extends Optimizer {
 
       tidy(() => {
         const newAccumulatedGrad =
-            this.rho.mul(accumulatedGrad)
+            this.rhoScalar.mul(accumulatedGrad)
                 .add(this.oneMinusRho.mul(gradient.square()));
 
-        const updates = accumulatedUpdate.add(this.epsilon)
+        const updates = accumulatedUpdate.add(this.epsilonScalar)
                             .sqrt()
-                            .div(accumulatedGrad.add(this.epsilon).sqrt())
+                            .div(accumulatedGrad.add(this.epsilonScalar).sqrt())
                             .mul(gradient);
 
         const newAccumulatedUpdate =
-            this.rho.mul(accumulatedUpdate)
+            this.rhoScalar.mul(accumulatedUpdate)
                 .add(this.oneMinusRho.mul(updates.square()));
 
         this.accumulatedGrads[variableName].assign(newAccumulatedGrad);
@@ -87,8 +93,8 @@ export class AdadeltaOptimizer extends Optimizer {
 
   dispose() {
     this.c.dispose();
-    this.epsilon.dispose();
-    this.rho.dispose();
+    this.epsilonScalar.dispose();
+    this.rhoScalar.dispose();
     this.oneMinusRho.dispose();
     if (this.accumulatedUpdates != null) {
       Object.keys(this.accumulatedUpdates)
@@ -97,4 +103,16 @@ export class AdadeltaOptimizer extends Optimizer {
           .forEach(name => this.accumulatedGrads[name].dispose());
     }
   }
+  getConfig(): ConfigDict {
+    return {
+      learningRate: this.learningRate,
+      rho: this.rho,
+      epsilon: this.epsilon
+    };
+  }
+  static fromConfig<T extends Serializable>(
+      cls: SerializableConstructor<T>, config: ConfigDict): T {
+    return new cls(config.learningRate, config.rho, config.epsilon);
+  }
 }
+SerializationMap.register(AdadeltaOptimizer);
diff --git a/src/optimizers/adadelta_optimizer_test.ts b/src/optimizers/adadelta_optimizer_test.ts
@@ -16,8 +16,8 @@
  */
 
 import * as tf from '../index';
-import {ALL_ENVS, expectArraysClose} from '../test_util';
 import {describeWithFlags} from '../jasmine_util';
+import {ALL_ENVS, expectArraysClose} from '../test_util';
 
 describeWithFlags('AdadeltaOptimizer', ALL_ENVS, () => {
   it('basic', () => {
@@ -75,4 +75,10 @@ describeWithFlags('AdadeltaOptimizer', ALL_ENVS, () => {
     // The only tensor remaining is the argument to variable().
     expect(tf.memory().numTensors).toBe(1);
   });
+  it('serialization round-trip', () => {
+    const originalOpt = tf.train.adadelta(0.1, 0.2, 2e-8);
+    const reserialized = tf.AdadeltaOptimizer.fromConfig(
+        tf.AdadeltaOptimizer, originalOpt.getConfig());
+    expect(reserialized.getConfig()).toEqual(originalOpt.getConfig());
+  });
 });
diff --git a/src/optimizers/adagrad_optimizer.ts b/src/optimizers/adagrad_optimizer.ts
@@ -18,13 +18,16 @@
 import {ENV} from '../environment';
 import {keep, tidy} from '../globals';
 import {fill, scalar} from '../ops/ops';
+// tslint:disable-next-line:max-line-length
+import {ConfigDict, Serializable, SerializableConstructor, SerializationMap} from '../serialization';
 import {Scalar} from '../tensor';
 import {NamedVariableMap} from '../types';
 
 import {Optimizer} from './optimizer';
 
 /** @doclink Optimizer */
 export class AdagradOptimizer extends Optimizer {
+  static className = 'AdagradOptimizer';
   private c: Scalar;
   private epsilon: Scalar;
 
@@ -73,4 +76,15 @@ export class AdagradOptimizer extends Optimizer {
           .forEach(name => this.accumulatedGrads[name].dispose());
     }
   }
+  getConfig(): ConfigDict {
+    return {
+      learningRate: this.learningRate,
+      initialAccumulatorValue: this.initialAccumulatorValue,
+    };
+  }
+  static fromConfig<T extends Serializable>(
+      cls: SerializableConstructor<T>, config: ConfigDict): T {
+    return new cls(config.learningRate, config.initialAccumulatorValue);
+  }
 }
+SerializationMap.register(AdagradOptimizer);
diff --git a/src/optimizers/adagrad_optimizer_test.ts b/src/optimizers/adagrad_optimizer_test.ts
@@ -16,8 +16,8 @@
  */
 
 import * as tf from '../index';
-import {ALL_ENVS, expectArraysClose} from '../test_util';
 import {describeWithFlags} from '../jasmine_util';
+import {ALL_ENVS, expectArraysClose} from '../test_util';
 
 describeWithFlags('AdagradOptimizer', ALL_ENVS, () => {
   it('basic', () => {
@@ -69,4 +69,10 @@ describeWithFlags('AdagradOptimizer', ALL_ENVS, () => {
     // The only tensor remaining is the argument to variable().
     expect(tf.memory().numTensors).toBe(1);
   });
+  it('serialization round-trip', () => {
+    const originalOpt = tf.train.adagrad(0.1, 0.2);
+    const reserialized = tf.AdagradOptimizer.fromConfig(
+        tf.AdagradOptimizer, originalOpt.getConfig());
+    expect(reserialized.getConfig()).toEqual(originalOpt.getConfig());
+  });
 });
diff --git a/src/optimizers/adam_optimizer.ts b/src/optimizers/adam_optimizer.ts
@@ -18,15 +18,19 @@
 import {ENV} from '../environment';
 import {keep, tidy} from '../globals';
 import {scalar, zerosLike} from '../ops/ops';
+// tslint:disable-next-line:max-line-length
+import {ConfigDict, Serializable, SerializableConstructor, SerializationMap} from '../serialization';
 import {Scalar, Variable} from '../tensor';
 import {NamedVariableMap} from '../types';
+
 import {Optimizer} from './optimizer';
 
 export class AdamOptimizer extends Optimizer {
+  static className = 'AdamOptimizer';
   private c: Scalar;
-  private eps: Scalar;
-  private beta1: Scalar;
-  private beta2: Scalar;
+  private epsScalar: Scalar;
+  private beta1Scalar: Scalar;
+  private beta2Scalar: Scalar;
   private accBeta1: Variable;
   private accBeta2: Variable;
   private oneMinusBeta1: Scalar;
@@ -37,14 +41,14 @@ export class AdamOptimizer extends Optimizer {
   private accumulatedSecondMoment: NamedVariableMap = {};
 
   constructor(
-      protected learningRate: number, beta1: number, beta2: number,
-      epsilon = 1e-8) {
+      protected learningRate: number, protected beta1: number,
+      protected beta2: number, protected epsilon = 1e-8) {
     super();
     this.c = keep(scalar(-learningRate));
-    this.eps = keep(scalar(epsilon));
+    this.epsScalar = keep(scalar(epsilon));
     // b1, b2 keep initial value of beta* hyperparameters.
-    this.beta1 = keep(scalar(beta1));
-    this.beta2 = keep(scalar(beta2));
+    this.beta1Scalar = keep(scalar(beta1));
+    this.beta2Scalar = keep(scalar(beta2));
     tidy(() => {
       // accB* will be updated by batch.
       this.accBeta1 = scalar(beta1).variable();
@@ -77,10 +81,10 @@ export class AdamOptimizer extends Optimizer {
         const firstMoment = this.accumulatedFirstMoment[variableName];
         const secondMoment = this.accumulatedSecondMoment[variableName];
 
-        const newFirstMoment =
-            this.beta1.mul(firstMoment).add(this.oneMinusBeta1.mul(gradient));
+        const newFirstMoment = this.beta1Scalar.mul(firstMoment)
+                                   .add(this.oneMinusBeta1.mul(gradient));
         const newSecondMoment =
-            this.beta2.mul(secondMoment)
+            this.beta2Scalar.mul(secondMoment)
                 .add(this.oneMinusBeta2.mul(gradient.square()));
 
         const biasCorrectedFirstMoment = newFirstMoment.div(oneMinusAccBeta1);
@@ -89,23 +93,24 @@ export class AdamOptimizer extends Optimizer {
         this.accumulatedFirstMoment[variableName].assign(newFirstMoment);
         this.accumulatedSecondMoment[variableName].assign(newSecondMoment);
 
-        const newValue = this.c
-                             .mul(biasCorrectedFirstMoment.div(this.eps.add(
-                                 biasCorrectedSecondMoment.sqrt())))
-                             .add(value);
+        const newValue =
+            this.c
+                .mul(biasCorrectedFirstMoment.div(
+                    this.epsScalar.add(biasCorrectedSecondMoment.sqrt())))
+                .add(value);
         value.assign(newValue);
       }
 
-      this.accBeta1.assign(this.accBeta1.mul(this.beta1));
-      this.accBeta2.assign(this.accBeta2.mul(this.beta2));
+      this.accBeta1.assign(this.accBeta1.mul(this.beta1Scalar));
+      this.accBeta2.assign(this.accBeta2.mul(this.beta2Scalar));
     });
   }
 
   dispose() {
     this.c.dispose();
-    this.eps.dispose();
-    this.beta1.dispose();
-    this.beta2.dispose();
+    this.epsScalar.dispose();
+    this.beta1Scalar.dispose();
+    this.beta2Scalar.dispose();
     this.accBeta1.dispose();
     this.accBeta2.dispose();
     this.oneMinusBeta1.dispose();
@@ -122,4 +127,18 @@ export class AdamOptimizer extends Optimizer {
           .forEach(name => this.accumulatedSecondMoment[name].dispose());
     }
   }
+  getConfig(): ConfigDict {
+    return {
+      learningRate: this.learningRate,
+      beta1: this.beta1,
+      beta2: this.beta2,
+      epsilon: this.epsilon,
+    };
+  }
+  static fromConfig<T extends Serializable>(
+      cls: SerializableConstructor<T>, config: ConfigDict): T {
+    return new cls(
+        config.learningRate, config.beta1, config.beta2, config.epsilon);
+  }
 }
+SerializationMap.register(AdamOptimizer);
diff --git a/src/optimizers/adam_optimizer_test.ts b/src/optimizers/adam_optimizer_test.ts
@@ -16,8 +16,8 @@
  */
 
 import * as tf from '../index';
-import {ALL_ENVS, expectArraysClose} from '../test_util';
 import {describeWithFlags} from '../jasmine_util';
+import {ALL_ENVS, expectArraysClose} from '../test_util';
 
 describeWithFlags('AdamOptimizer', ALL_ENVS, () => {
   it('basic', () => {
@@ -79,4 +79,10 @@ describeWithFlags('AdamOptimizer', ALL_ENVS, () => {
     // The only tensor remaining should be the argument to variable().
     expect(tf.memory().numTensors).toBe(1);
   });
+  it('serialization round-trip', () => {
+    const originalOpt = tf.train.adam(0.1, 0.2, 0.3, 2e-8);
+    const reserialized =
+        tf.AdamOptimizer.fromConfig(tf.AdamOptimizer, originalOpt.getConfig());
+    expect(reserialized.getConfig()).toEqual(originalOpt.getConfig());
+  });
 });