[io] Add field user-defined metadata; Tweak API for tf.io.fromMemory() (#1864)

FEATURE

- This is the first PR for adding support for user-defined metadata in model artifacts. Design doc has been circulated and discussed.
- Add the field `userDefinedMetadata` to `ModelArtifacts` and `ModelJSON`.
- Deprecate the old API of `tf.io.fromMemory()` which consisted of multiple arguments. The arguments are consolidated into on in the new API.
- Add unit tests.

Towards tensorflow/tfjs#1596

Author: caisq

src/io/passthrough.ts

@@ -22,27 +22,10 @@
 import {IOHandler, ModelArtifacts, SaveResult, TrainingConfig, WeightsManifestEntry} from './types';
 
 class PassthroughLoader implements IOHandler {
-  constructor(
-      private readonly modelTopology?: {}|ArrayBuffer,
-      private readonly weightSpecs?: WeightsManifestEntry[],
-      private readonly weightData?: ArrayBuffer,
-      private readonly trainingConfig?: TrainingConfig) {}
+  constructor(private readonly modelArtifacts?: ModelArtifacts) {}
 
   async load(): Promise<ModelArtifacts> {
-    let result = {};
-    if (this.modelTopology != null) {
-      result = {modelTopology: this.modelTopology, ...result};
-    }
-    if (this.weightSpecs != null && this.weightSpecs.length > 0) {
-      result = {weightSpecs: this.weightSpecs, ...result};
-    }
-    if (this.weightData != null && this.weightData.byteLength > 0) {
-      result = {weightData: this.weightData, ...result};
-    }
-    if (this.trainingConfig != null) {
-      result = {trainingConfig: this.trainingConfig, ...result};
-    }
-    return result;
+    return this.modelArtifacts;
   }
 }
 
@@ -67,7 +50,7 @@ class PassthroughSaver implements IOHandler {
  *     modelTopology, weightSpecs, weightData));
  * ```
  *
- * @param modelTopology a object containing model topology (i.e., parsed from
+ * @param modelArtifacts a object containing model topology (i.e., parsed from
  *   the JSON format).
  * @param weightSpecs An array of `WeightsManifestEntry` objects describing the
  *   names, shapes, types, and quantization of the weight data.
@@ -78,13 +61,39 @@ class PassthroughSaver implements IOHandler {
  * @returns A passthrough `IOHandler` that simply loads the provided data.
  */
 export function fromMemory(
-    modelTopology: {}, weightSpecs?: WeightsManifestEntry[],
+    modelArtifacts: {}|ModelArtifacts, weightSpecs?: WeightsManifestEntry[],
     weightData?: ArrayBuffer, trainingConfig?: TrainingConfig): IOHandler {
-  // TODO(cais): The arguments should probably be consolidated into a single
-  // object, with proper deprecation process. Even though this function isn't
-  // documented, it is public and being used by some downstream libraries.
-  return new PassthroughLoader(
-      modelTopology, weightSpecs, weightData, trainingConfig);
+  if (arguments.length === 1) {
+    const isModelArtifacts =
+        (modelArtifacts as ModelArtifacts).modelTopology != null ||
+        (modelArtifacts as ModelArtifacts).weightSpecs != null;
+    if (isModelArtifacts) {
+      return new PassthroughLoader(modelArtifacts as ModelArtifacts);
+    } else {
+      // Legacy support: with only modelTopology.
+      // TODO(cais): Remove this deprecated API.
+      console.warn(
+          'Please call tf.io.fromMemory() with only one argument. ' +
+          'The argument should be of type ModelArtifacts. ' +
+          'The multi-argument signature of tf.io.fromMemory() has been ' +
+          'deprecated and will be removed in a future release.');
+      return new PassthroughLoader({modelTopology: modelArtifacts as {}});
+    }
+  } else {
+    // Legacy support.
+    // TODO(cais): Remove this deprecated API.
+    console.warn(
+        'Please call tf.io.fromMemory() with only one argument. ' +
+        'The argument should be of type ModelArtifacts. ' +
+        'The multi-argument signature of tf.io.fromMemory() has been ' +
+        'deprecated and will be removed in a future release.');
+    return new PassthroughLoader({
+      modelTopology: modelArtifacts as {},
+      weightSpecs,
+      weightData,
+      trainingConfig
+    });
+  }
 }
 
 /**

src/io/passthrough_test.ts

@@ -114,20 +114,61 @@ describeWithFlags('Passthrough Saver', BROWSER_ENVS, () => {
 });
 
 describeWithFlags('Passthrough Loader', BROWSER_ENVS, () => {
+  it('load topology and weights: legacy signature', async () => {
+    const passthroughHandler = tf.io.fromMemory(
+        modelTopology1, weightSpecs1, weightData1);
+    const modelArtifacts = await passthroughHandler.load();
+    expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
+    expect(modelArtifacts.weightSpecs).toEqual(weightSpecs1);
+    expect(modelArtifacts.weightData).toEqual(weightData1);
+    expect(modelArtifacts.userDefinedMetadata).toEqual(undefined);
+  });
+
   it('load topology and weights', async () => {
-    const passthroughHandler =
-        tf.io.fromMemory(modelTopology1, weightSpecs1, weightData1);
+    const passthroughHandler = tf.io.fromMemory({
+      modelTopology: modelTopology1,
+      weightSpecs: weightSpecs1,
+      weightData: weightData1
+    });
     const modelArtifacts = await passthroughHandler.load();
     expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
     expect(modelArtifacts.weightSpecs).toEqual(weightSpecs1);
     expect(modelArtifacts.weightData).toEqual(weightData1);
+    expect(modelArtifacts.userDefinedMetadata).toEqual(undefined);
   });
 
-  it('load model topology only', async () => {
+  it('load model topology only: legacy signature', async () => {
     const passthroughHandler = tf.io.fromMemory(modelTopology1);
     const modelArtifacts = await passthroughHandler.load();
     expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
     expect(modelArtifacts.weightSpecs).toEqual(undefined);
     expect(modelArtifacts.weightData).toEqual(undefined);
+    expect(modelArtifacts.userDefinedMetadata).toEqual(undefined);
+  });
+
+  it('load model topology only', async () => {
+    const passthroughHandler = tf.io.fromMemory({
+      modelTopology: modelTopology1
+    });
+    const modelArtifacts = await passthroughHandler.load();
+    expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
+    expect(modelArtifacts.weightSpecs).toEqual(undefined);
+    expect(modelArtifacts.weightData).toEqual(undefined);
+    expect(modelArtifacts.userDefinedMetadata).toEqual(undefined);
+  });
+
+  it('load topology, weights, and user-defined metadata', async () => {
+    const userDefinedMetadata: {} = {'fooField': 'fooValue'};
+    const passthroughHandler = tf.io.fromMemory({
+      modelTopology: modelTopology1,
+      weightSpecs: weightSpecs1,
+      weightData: weightData1,
+      userDefinedMetadata
+    });
+    const modelArtifacts = await passthroughHandler.load();
+    expect(modelArtifacts.modelTopology).toEqual(modelTopology1);
+    expect(modelArtifacts.weightSpecs).toEqual(weightSpecs1);
+    expect(modelArtifacts.weightData).toEqual(weightData1);
+    expect(modelArtifacts.userDefinedMetadata).toEqual(userDefinedMetadata);
   });
 });

src/io/types.ts

@@ -275,6 +275,11 @@ export declare interface ModelArtifacts {
    * `tf.LayersModel` instance.)
    */
   convertedBy?: string|null;
+
+  /**
+   * User-defined metadata about the model.
+   */
+  userDefinedMetadata?: {};
 }
 
 /**
@@ -330,6 +335,11 @@ export declare interface ModelJSON {
    * `tf.LayersModel` instance.)
    */
   convertedBy?: string|null;
+
+  /**
+   * User-defined metadata about the model.
+   */
+  userDefinedMetadata?: {};
 }
 
 /**