Add support for user-defined metadata; Simplify fromMemory() usage (#571)

FEATURE

- Add two public methods to `tf.LayersModel` and `tf.Sequential`
  - `setUserDefinedMetadata()`
  - `getUserDefinedMetadata()`
  These methods allow user to set and get custom metadata about the model.
- Any set user metadata is serialized together with the model.
- User-defined metadata are also deserialized (i.e., loaded) together with the topology and/or weights of the model.
- The user-defined metadata is required to be a plain JSON object. This PR implements the enforcement mechanisms that will throw errors if this condition is not met.
- During serialization and deserializatoin, the size of the JSON object will be checked. If it is greater than 1 MB in length, a warning will be thrown.

Towards #1596

Author: caisq

src/engine/training.ts

@@ -27,6 +27,7 @@ import * as losses from '../losses';
 import * as Metrics from '../metrics';
 import * as optimizers from '../optimizers';
 import {LossOrMetricFn} from '../types';
+import {checkUserDefinedMetadata} from '../user_defined_metadata';
 import {count, pyListRepeat, singletonOrArray, toCamelCase, toSnakeCase, unique} from '../utils/generic_utils';
 import {printSummary} from '../utils/layer_utils';
 import {range} from '../utils/math_utils';
@@ -506,6 +507,9 @@ export class LayersModel extends Container implements tfc.InferenceModel {
   //   implicit "knowledge" of the outputs it depends on.
   metricsTensors: Array<[LossOrMetricFn, number]>;
 
+  // User defind metadata (if any).
+  private userDefinedMetadata: {};
+
   constructor(args: ContainerArgs) {
     super(args);
     this.isTraining = false;
@@ -837,8 +841,8 @@ export class LayersModel extends Container implements tfc.InferenceModel {
     // TODO(cais): Standardize `config.sampleWeights` as well.
     // Validate user data.
     const checkBatchAxis = true;
-    const standardizedOuts = this.standardizeUserDataXY(
-        x, y, checkBatchAxis, batchSize);
+    const standardizedOuts =
+        this.standardizeUserDataXY(x, y, checkBatchAxis, batchSize);
     try {
       // TODO(cais): If uses `useLearningPhase`, set the corresponding element
       // of the input to 0.
@@ -1136,10 +1140,9 @@ export class LayersModel extends Container implements tfc.InferenceModel {
   }
 
   protected standardizeUserDataXY(
-    x: Tensor|Tensor[]|{[inputName: string]: Tensor},
-    y: Tensor|Tensor[]|{[inputName: string]: Tensor},
-    checkBatchAxis = true,
-    batchSize?: number): [Tensor[], Tensor[]] {
+      x: Tensor|Tensor[]|{[inputName: string]: Tensor},
+      y: Tensor|Tensor[]|{[inputName: string]: Tensor}, checkBatchAxis = true,
+      batchSize?: number): [Tensor[], Tensor[]] {
     // TODO(cais): Add sampleWeight, classWeight
     if (this.optimizer_ == null) {
       throw new RuntimeError(
@@ -1348,8 +1351,9 @@ export class LayersModel extends Container implements tfc.InferenceModel {
           } else {
             const metric = this.metricsTensors[i][0];
             const outputIndex = this.metricsTensors[i][1];
-            weightedMetric = tfc.mean(
-                metric(targets[outputIndex], outputs[outputIndex])) as Scalar;
+            weightedMetric =
+                tfc.mean(metric(targets[outputIndex], outputs[outputIndex])) as
+                Scalar;
           }
 
           tfc.keep(weightedMetric);
@@ -1547,7 +1551,7 @@ export class LayersModel extends Container implements tfc.InferenceModel {
    * @returns A `NamedTensorMap` mapping original weight names (i.e.,
    *   non-uniqueified weight names) to their values.
    */
-  protected getNamedWeights(config?: io.SaveConfig): NamedTensor [] {
+  protected getNamedWeights(config?: io.SaveConfig): NamedTensor[] {
     const namedWeights: NamedTensor[] = [];
 
     const trainableOnly = config != null && config.trainableOnly;
@@ -1642,15 +1646,15 @@ export class LayersModel extends Container implements tfc.InferenceModel {
     } else {
       const outputNames = Object.keys(this.loss);
       lossNames = {} as {[outputName: string]: LossIdentifier};
-      const losses = this.loss as {[outputName: string]: LossOrMetricFn|string};
+      const losses =
+          this.loss as {[outputName: string]: LossOrMetricFn | string};
       for (const outputName of outputNames) {
         if (typeof losses[outputName] === 'string') {
           lossNames[outputName] =
               toSnakeCase(losses[outputName] as string) as LossIdentifier;
         } else {
           throw new Error('Serialization of non-string loss is not supported.');
         }
-
       }
     }
     return lossNames;
@@ -1722,8 +1726,8 @@ export class LayersModel extends Container implements tfc.InferenceModel {
     } else if (trainingConfig.metrics != null) {
       metrics = {} as {[outputName: string]: MetricsIdentifier};
       for (const key in trainingConfig.metrics) {
-        metrics[key] = toCamelCase(trainingConfig.metrics[key]) as
-            MetricsIdentifier;
+        metrics[key] =
+            toCamelCase(trainingConfig.metrics[key]) as MetricsIdentifier;
       }
     }
 
@@ -1856,9 +1860,44 @@ export class LayersModel extends Container implements tfc.InferenceModel {
           [weightDataAndSpecs.data, optimizerWeightData]);
     }
 
+    if (this.userDefinedMetadata != null) {
+      // Check serialized size of user-defined metadata.
+      const checkSize = true;
+      checkUserDefinedMetadata(this.userDefinedMetadata, this.name, checkSize);
+      modelArtifacts.userDefinedMetadata = this.userDefinedMetadata;
+    }
+
     modelArtifacts.weightData = weightDataAndSpecs.data;
     modelArtifacts.weightSpecs = weightDataAndSpecs.specs;
     return handlerOrURL.save(modelArtifacts);
   }
+
+  /**
+   * Set user-defined metadata.
+   *
+   * The set metadata will be serialized together with the topology
+   * and weights of the model during `save()` calls.
+   *
+   * @param setUserDefinedMetadata
+   */
+  setUserDefinedMetadata(userDefinedMetadata: {}): void {
+    checkUserDefinedMetadata(userDefinedMetadata, this.name);
+    this.userDefinedMetadata = userDefinedMetadata;
+  }
+
+  /**
+   * Get user-defined metadata.
+   *
+   * The metadata is supplied via one of the two routes:
+   *   1. By calling `setUserDefinedMetadata()`.
+   *   2. Loaded during model loading (if the model is constructed
+   *      via `tf.loadLayersModel()`.)
+   *
+   * If no user-defined metadata is available from either of the
+   * two routes, this function will return `undefined`.
+   */
+  getUserDefinedMetadata(): {} {
+    return this.userDefinedMetadata;
+  }
 }
 serialization.registerClass(LayersModel);

src/model_save_test.ts

@@ -290,9 +290,7 @@ describeMathGPU('Save-load round trips', () => {
 
     const getInitSpy = spyOn(initializers, 'getInitializer').and.callThrough();
     const gramSchmidtSpy = spyOn(linalg, 'gramSchmidt').and.callThrough();
-    const modelPrime = await tfl.loadLayersModel(io.fromMemory(
-        savedArtifacts.modelTopology, savedArtifacts.weightSpecs,
-        savedArtifacts.weightData));
+    const modelPrime = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     const weightsPrime = modelPrime.getWeights();
     expect(weightsPrime.length).toEqual(weights.length);
     for (let i = 0; i < weights.length; ++i) {
@@ -321,9 +319,7 @@ describeMathGPU('Save-load round trips', () => {
 
     const getInitSpy = spyOn(initializers, 'getInitializer').and.callThrough();
     const gramSchmidtSpy = spyOn(linalg, 'gramSchmidt').and.callThrough();
-    const modelPrime = await tfl.loadLayersModel(io.fromMemory(
-        savedArtifacts.modelTopology, savedArtifacts.weightSpecs,
-        savedArtifacts.weightData));
+    const modelPrime = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     const weightsPrime = modelPrime.getWeights();
     expect(weightsPrime.length).toEqual(weights.length);
     for (let i = 0; i < weights.length; ++i) {
@@ -357,9 +353,7 @@ describeMathGPU('Save-load round trips', () => {
 
     const getInitSpy = spyOn(initializers, 'getInitializer').and.callThrough();
     const gramSchmidtSpy = spyOn(linalg, 'gramSchmidt').and.callThrough();
-    const modelPrime = await tfl.loadLayersModel(io.fromMemory(
-        savedArtifacts.modelTopology, savedArtifacts.weightSpecs,
-        savedArtifacts.weightData));
+    const modelPrime = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     const weightsPrime = modelPrime.getWeights();
     expect(weightsPrime.length).toEqual(weights.length);
     for (let i = 0; i < weights.length; ++i) {
@@ -395,9 +389,7 @@ describeMathGPU('Save-load round trips', () => {
 
     const getInitSpy = spyOn(initializers, 'getInitializer').and.callThrough();
     const gramSchmidtSpy = spyOn(linalg, 'gramSchmidt').and.callThrough();
-    const modelPrime = await tfl.loadLayersModel(io.fromMemory(
-        savedArtifacts.modelTopology, savedArtifacts.weightSpecs,
-        savedArtifacts.weightData));
+    const modelPrime = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     const weightsPrime = modelPrime.getWeights();
     expect(weightsPrime.length).toEqual(weights.length);
     for (let i = 0; i < weights.length; ++i) {
@@ -452,10 +444,7 @@ describeMathGPU('Save-load round trips', () => {
     const gramSchmidtSpy = spyOn(linalg, 'gramSchmidt').and.callThrough();
     const strict = false;
     const modelPrime = await tfl.loadLayersModel(
-        io.fromMemory(
-            savedArtifacts.modelTopology, savedArtifacts.weightSpecs,
-            savedArtifacts.weightData),
-        {strict});
+        io.fromMemory(savedArtifacts), {strict});
     const weightsPrime = modelPrime.getWeights();
     expect(weightsPrime.length).toEqual(weights.length);
     expectTensorsClose(weightsPrime[0], weights[0]);
@@ -486,10 +475,7 @@ describeMathGPU('Save-load round trips', () => {
     const gramSchmidtSpy = spyOn(linalg, 'gramSchmidt').and.callThrough();
     const strict = false;
     const modelPrime = await tfl.loadLayersModel(
-        io.fromMemory(
-            savedArtifacts.modelTopology, savedArtifacts.weightSpecs,
-            savedArtifacts.weightData),
-        {strict});
+        io.fromMemory(savedArtifacts), {strict});
     const weightsPrime = modelPrime.getWeights();
     expect(weightsPrime.length).toEqual(weights.length);
     expectTensorsClose(weightsPrime[0], weights[0]);

src/models.ts

@@ -306,6 +306,9 @@ export async function loadLayersModelFromIOHandler(
   if (trainingConfig != null) {
     model.loadTrainingConfig(trainingConfig);
   }
+  if (artifacts.userDefinedMetadata != null) {
+    model.setUserDefinedMetadata(artifacts.userDefinedMetadata);
+  }
 
   // If weightData is present, load the weights into the model.
   if (artifacts.weightData != null) {

src/models_test.ts

@@ -1237,7 +1237,6 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual('mean_squared_error');
 
@@ -1254,8 +1253,7 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
     expect(weightData.byteLength).toEqual(4 * 8 + 4 * 1 + 4);
 
     // Load the model back, with the optimizer.
-    const model2 = await tfl.loadLayersModel(
-        io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     expect(model2.optimizer.getConfig()['learningRate']).toEqual(learningRate);
 
     const optimizer1Weights = await model1.optimizer.getWeights();
@@ -1295,7 +1293,6 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual('mean_squared_error');
 
@@ -1317,8 +1314,7 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
     expect(weightData.byteLength).toEqual(4 + 4 * 8 * 3 + 4 * 1 * 3);
 
     // Load the model back, with the optimizer.
-    const model2 = await tfl.loadLayersModel(
-        io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     expect(model2.optimizer.getConfig()['learningRate']).toEqual(learningRate);
     expect(model2.optimizer.getConfig()['decay']).toEqual(decay);
 
@@ -1358,7 +1354,6 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual('mean_squared_error');
 
@@ -1380,8 +1375,7 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
     expect(weightData.byteLength).toEqual(4 + 4 * 8 * 3 + 4 * 1 * 3);
 
     // Load the model back, with the optimizer.
-    const model2 = await tfl.loadLayersModel(
-        io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     expect(model2.optimizer.getConfig()['learningRate']).toEqual(learningRate);
 
     const optimizer1Weights = await model1.optimizer.getWeights();
@@ -1421,7 +1415,6 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual('mean_squared_error');
 
@@ -1441,8 +1434,7 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
     expect(weightData.byteLength).toEqual(4 + 4 * 8 * 2 + 4 * 1 * 2);
 
     // Load the model back, with the optimizer.
-    const model2 = await tfl.loadLayersModel(
-        io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     expect(model2.optimizer.getConfig()['learningRate']).toEqual(learningRate);
     expect(model2.optimizer.getConfig()['initialAccumulatorValue'])
         .toEqual(initialAccumulatorValue);
@@ -1485,7 +1477,6 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual('mean_squared_error');
 
@@ -1507,8 +1498,7 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
     expect(weightData.byteLength).toEqual(4 + 4 * 8 * 3 + 4 * 1 * 3);
 
     // Load the model back, with the optimizer.
-    const model2 = await tfl.loadLayersModel(
-        io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     expect(model2.optimizer.getConfig()['learningRate']).toEqual(learningRate);
     expect(model2.optimizer.getConfig()['beta1']).toEqual(beta1);
     expect(model2.optimizer.getConfig()['beta2']).toEqual(beta2);
@@ -1547,7 +1537,6 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual('mean_squared_error');
 
@@ -1569,8 +1558,7 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
     expect(weightData.byteLength).toEqual(4 + 4 * 8 * 3 + 4 * 1 * 3);
 
     // Load the model back, with the optimizer.
-    const model2 = await tfl.loadLayersModel(
-        io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     expect(model2.optimizer.getConfig()['learningRate']).toEqual(1e-3);
 
     const optimizer1Weights = await model1.optimizer.getWeights();
@@ -1610,7 +1598,6 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual('mean_squared_error');
 
@@ -1630,8 +1617,7 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
     expect(weightData.byteLength).toEqual(4 + 4 * 8 * 2 + 4 * 1 * 2);
 
     // Load the model back, with the optimizer.
-    const model2 = await tfl.loadLayersModel(
-        io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     expect(model2.optimizer.getConfig()['learningRate']).toEqual(learningRate);
 
     const optimizer1Weights = await model1.optimizer.getWeights();
@@ -1675,15 +1661,11 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual('categorical_crossentropy');
     expect(trainingConfig['metrics']).toEqual(['acc']);
 
-    const weightSpecs = savedArtifacts.weightSpecs;
-    const weightData = savedArtifacts.weightData;
-    const model2 = await tfl.loadLayersModel(
-      io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
     h = await model2.fit(xs, ys, {epochs: 1});
     expect(h.history.loss.length).toEqual(1);
     expect(h.history.loss[0]).toBeCloseTo(1.086648);
@@ -1711,15 +1693,11 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual('categorical_crossentropy');
     expect(trainingConfig['metrics']).toEqual(['acc']);
 
-    const weightSpecs = savedArtifacts.weightSpecs;
-    const weightData = savedArtifacts.weightData;
-    const model2 = await tfl.loadLayersModel(
-      io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
 
     const xs = ones([4, 8]);
     const ys = tensor2d([[0, 0, 1], [0, 0, 1], [0, 0, 1], [0, 0, 1]]);
@@ -1769,16 +1747,12 @@ describeMathCPUAndGPU('Saving+loading model with optimizer', () => {
           return null;
         }), {includeOptimizer: true});
 
-    const modelTopology = savedArtifacts.modelTopology as ConfigDict;
     const trainingConfig = savedArtifacts.trainingConfig;
     expect(trainingConfig['loss']).toEqual(
         ['categorical_crossentropy', 'binary_crossentropy']);
     expect(trainingConfig['metrics']).toEqual(['acc']);
 
-    const weightSpecs = savedArtifacts.weightSpecs;
-    const weightData = savedArtifacts.weightData;
-    const model2 = await tfl.loadLayersModel(
-      io.fromMemory(modelTopology, weightSpecs, weightData, trainingConfig));
+    const model2 = await tfl.loadLayersModel(io.fromMemory(savedArtifacts));
 
     h = await model2.fit(xs, [ys1, ys2], {epochs: 1});
     expect(h.history.loss.length).toEqual(1);