Add fields to io.ModelArtifacts and surrounding save/load logic (tensorflow#1596)

- Add new fields to `ModelArtifacts` interface
  - format
  - generatedBy
  - convertedBy
- Add the logic to save and load these new fields, with tests, to:
  - Local Storage
  - HTTP
  - Browser file downloads
  - IndexedDB (test only; no non-test code change needed)

Towards: tensorflow/tfjs#1285
Towards: tensorflow/tfjs#1286

FEATURE

Author: caisq

src/io/browser_files.ts

@@ -73,6 +73,9 @@ export class BrowserDownloads implements IOHandler {
       }];
       const modelTopologyAndWeightManifest = {
         modelTopology: modelArtifacts.modelTopology,
+        format: modelArtifacts.format,
+        generatedBy: modelArtifacts.generatedBy,
+        convertedBy: modelArtifacts.convertedBy,
         weightsManifest
       };
       const modelTopologyAndWeightManifestURL =

src/io/browser_files_test.ts

@@ -74,6 +74,9 @@ const artifacts1: tf.io.ModelArtifacts = {
   modelTopology: modelTopology1,
   weightSpecs: weightSpecs1,
   weightData: weightData1,
+  format: 'layers-model',
+  generatedBy: 'TensorFlow.js v0.0.0',
+  convertedBy: null
 };
 
 describeWithFlags('browserDownloads', BROWSER_ENVS, () => {
@@ -127,6 +130,10 @@ describeWithFlags('browserDownloads', BROWSER_ENVS, () => {
         JSON.parse(await jsonContent.text());
     expect(modelTopologyAndWeightsManifest.modelTopology)
         .toEqual(modelTopology1);
+    expect(modelTopologyAndWeightsManifest.format).toEqual('layers-model');
+    expect(modelTopologyAndWeightsManifest.generatedBy)
+        .toEqual('TensorFlow.js v0.0.0');
+    expect(modelTopologyAndWeightsManifest.convertedBy).toEqual(null);
     const weightsManifest = modelTopologyAndWeightsManifest.weightsManifest as
         WeightsManifestConfig;
     expect(weightsManifest.length).toEqual(1);

src/io/browser_http.ts

@@ -114,6 +114,9 @@ export class BrowserHTTPRequest implements IOHandler {
     }];
     const modelTopologyAndWeightManifest = {
       modelTopology: modelArtifacts.modelTopology,
+      format: modelArtifacts.format,
+      generatedBy: modelArtifacts.generatedBy,
+      convertedBy: modelArtifacts.convertedBy,
       weightsManifest
     };
 

src/io/browser_http_test.ts

@@ -192,6 +192,9 @@ describeWithFlags('browserHTTPRequest-save', CHROME_ENVS, () => {
     modelTopology: modelTopology1,
     weightSpecs: weightSpecs1,
     weightData: weightData1,
+    format: 'layers-model',
+    generatedBy: 'TensorFlow.js v0.0.0',
+    convertedBy: null
   };
 
   let requestInits: RequestInit[] = [];
@@ -338,6 +341,9 @@ describeWithFlags('browserHTTPRequest-save', CHROME_ENVS, () => {
           jsonFileReader.onload = (event: Event) => {
             // tslint:disable-next-line:no-any
             const modelJSON = JSON.parse((event.target as any).result);
+            expect(modelJSON.format).toEqual('layers-model');
+            expect(modelJSON.generatedBy).toEqual('TensorFlow.js v0.0.0');
+            expect(modelJSON.convertedBy).toEqual(null);
             expect(modelJSON.modelTopology).toEqual(modelTopology1);
             expect(modelJSON.weightsManifest.length).toEqual(1);
             expect(modelJSON.weightsManifest[0].weights).toEqual(weightSpecs1);

src/io/indexed_db_test.ts

@@ -75,6 +75,9 @@ describeWithFlags('IndexedDB', BROWSER_ENVS, () => {
     modelTopology: modelTopology1,
     weightSpecs: weightSpecs1,
     weightData: weightData1,
+    format: 'layers-model',
+    generatedBy: 'TensorFlow.js v0.0.0',
+    convertedBy: null
   };
 
   const weightSpecs2: tf.io.WeightsManifestEntry[] = [
@@ -113,6 +116,9 @@ describeWithFlags('IndexedDB', BROWSER_ENVS, () => {
     const loadedArtifacts = await handler.load();
     expect(loadedArtifacts.modelTopology).toEqual(modelTopology1);
     expect(loadedArtifacts.weightSpecs).toEqual(weightSpecs1);
+    expect(loadedArtifacts.format).toEqual('layers-model');
+    expect(loadedArtifacts.generatedBy).toEqual('TensorFlow.js v0.0.0');
+    expect(loadedArtifacts.convertedBy).toEqual(null);
     expectArrayBuffersEqual(loadedArtifacts.weightData, weightData1);
   });
 

src/io/io.ts

@@ -25,7 +25,7 @@ import {browserHTTPRequest, isHTTPScheme} from './browser_http';
 import {concatenateArrayBuffers, decodeWeights, encodeWeights, getModelArtifactsInfoForJSON} from './io_utils';
 import {fromMemory, withSaveHandler} from './passthrough';
 import {getLoadHandlers, getSaveHandlers, registerLoadRouter, registerSaveRouter} from './router_registry';
-import {IOHandler, LoadHandler, LoadOptions, ModelArtifacts, ModelStoreManager, OnProgressCallback, SaveConfig, SaveHandler, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
+import {IOHandler, LoadHandler, LoadOptions, ModelArtifacts, ModelStoreManager, ModelFormat, OnProgressCallback, SaveConfig, SaveHandler, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
 import {loadWeights, weightsLoaderFactory} from './weights_loader';
 
 export {copyModel, listModels, moveModel, removeModel} from './model_management';
@@ -45,6 +45,7 @@ export {
   LoadOptions,
   loadWeights,
   ModelArtifacts,
+  ModelFormat,
   ModelStoreManager,
   OnProgressCallback,
   registerLoadRouter,

src/io/local_storage.ts

@@ -20,14 +20,15 @@ import {assert} from '../util';
 import {arrayBufferToBase64String, base64StringToArrayBuffer, getModelArtifactsInfoForJSON} from './io_utils';
 import {ModelStoreManagerRegistry} from './model_management';
 import {IORouter, IORouterRegistry} from './router_registry';
-import {IOHandler, ModelArtifacts, ModelArtifactsInfo, ModelStoreManager, SaveResult} from './types';
+import {IOHandler, ModelArtifacts, ModelArtifactsInfo, ModelStoreManager, ModelFormat, SaveResult} from './types';
 
 const PATH_SEPARATOR = '/';
 const PATH_PREFIX = 'tensorflowjs_models';
 const INFO_SUFFIX = 'info';
 const MODEL_TOPOLOGY_SUFFIX = 'model_topology';
 const WEIGHT_SPECS_SUFFIX = 'weight_specs';
 const WEIGHT_DATA_SUFFIX = 'weight_data';
+const MODEL_METADATA_SUFFIX = 'model_metadata';
 
 /**
  * Purge all tensorflow.js-saved model artifacts from local storage.
@@ -57,12 +58,15 @@ export function purgeLocalStorageArtifacts(): string[] {
 }
 
 function getModelKeys(path: string):
-    {info: string, topology: string, weightSpecs: string, weightData: string} {
+    {info: string, topology: string, weightSpecs: string, weightData: string,
+     modelMetadata: string} {
   return {
     info: [PATH_PREFIX, path, INFO_SUFFIX].join(PATH_SEPARATOR),
     topology: [PATH_PREFIX, path, MODEL_TOPOLOGY_SUFFIX].join(PATH_SEPARATOR),
     weightSpecs: [PATH_PREFIX, path, WEIGHT_SPECS_SUFFIX].join(PATH_SEPARATOR),
-    weightData: [PATH_PREFIX, path, WEIGHT_DATA_SUFFIX].join(PATH_SEPARATOR)
+    weightData: [PATH_PREFIX, path, WEIGHT_DATA_SUFFIX].join(PATH_SEPARATOR),
+    modelMetadata: [
+        PATH_PREFIX, path, MODEL_METADATA_SUFFIX].join(PATH_SEPARATOR)
   };
 }
 
@@ -146,6 +150,11 @@ export class BrowserLocalStorage implements IOHandler {
         this.LS.setItem(
             this.keys.weightData,
             arrayBufferToBase64String(modelArtifacts.weightData));
+        this.LS.setItem(this.keys.modelMetadata, JSON.stringify({
+          format: modelArtifacts.format,
+          generatedBy: modelArtifacts.generatedBy,
+          convertedBy: modelArtifacts.convertedBy
+        }));
 
         return {modelArtifactsInfo};
       } catch (err) {
@@ -206,6 +215,19 @@ export class BrowserLocalStorage implements IOHandler {
     }
     out.weightSpecs = weightSpecs;
 
+    // Load meta-data fields.
+    const metadataString = this.LS.getItem(this.keys.modelMetadata);
+    if (metadataString != null) {
+      const metadata = JSON.parse(metadataString) as {
+        format: string,
+        generatedBy: string,
+        convertedBy: string
+      };
+      out.format = metadata.format as ModelFormat;
+      out.generatedBy = metadata['generatedBy'];
+      out.convertedBy = metadata['convertedBy'];
+    }
+
     // Load weight data.
     const weightDataBase64 = this.LS.getItem(this.keys.weightData);
     if (weightDataBase64 == null) {

src/io/local_storage_test.ts

@@ -68,10 +68,20 @@ describeWithFlags('LocalStorage', BROWSER_ENVS, () => {
     }
   ];
   const weightData1 = new ArrayBuffer(16);
+
   const artifacts1: tf.io.ModelArtifacts = {
     modelTopology: modelTopology1,
     weightSpecs: weightSpecs1,
     weightData: weightData1,
+    format: 'layers-model',
+    generatedBy: 'TensorFlow.js v0.0.0',
+    convertedBy: null
+  };
+
+  const artifactsV0: tf.io.ModelArtifacts = {
+    modelTopology: modelTopology1,
+    weightSpecs: weightSpecs1,
+    weightData: weightData1
   };
 
   function findOverflowingByteSize(): number {
@@ -162,6 +172,23 @@ describeWithFlags('LocalStorage', BROWSER_ENVS, () => {
     expect(loaded.modelTopology).toEqual(modelTopology1);
     expect(loaded.weightSpecs).toEqual(weightSpecs1);
     expect(loaded.weightData).toEqual(weightData1);
+    expect(loaded.format).toEqual('layers-model');
+    expect(loaded.generatedBy).toEqual('TensorFlow.js v0.0.0');
+    expect(loaded.convertedBy).toEqual(null);
+  });
+
+  it('Save-load round trip succeeds: v0 format', async () => {
+    const handler1 = tf.io.getSaveHandlers('localstorage://FooModel')[0];
+
+    await handler1.save(artifactsV0);
+    const handler2 = tf.io.getLoadHandlers('localstorage://FooModel')[0];
+    const loaded = await handler2.load();
+    expect(loaded.modelTopology).toEqual(modelTopology1);
+    expect(loaded.weightSpecs).toEqual(weightSpecs1);
+    expect(loaded.weightData).toEqual(weightData1);
+    expect(loaded.format).toEqual(undefined);
+    expect(loaded.generatedBy).toEqual(undefined);
+    expect(loaded.convertedBy).toEqual(undefined);
   });
 
   it('Loading nonexistent model fails.', done => {

src/io/types.ts

@@ -156,6 +156,8 @@ export declare interface ModelArtifactsInfo {
   weightDataBytes?: number;
 }
 
+export declare type ModelFormat = 'graph-model'|'layers-model';
+
 /**
  * The serialized artifacts of a model, including topology and weights.
  *
@@ -185,8 +187,38 @@ export declare interface ModelArtifacts {
    * by `weightSpecs`.
    */
   weightData?: ArrayBuffer;
+
+  /**
+   * Hard-coded format name for models saved from TensorFlow.js or converted
+   * by TensorFlow.js Converter.
+   */
+  format?: ModelFormat;
+
+  /**
+   * What library is responsible for originally generating this artifact.
+   *
+   * Used for debugging purposes. E.g., 'TensorFlow.js v1.0.0'.
+   */
+  generatedBy?: string;
+
+  /**
+   * What library or tool is responsible for converting the original model
+   * to this format, applicable only if the model is output by a converter.
+   *
+   * Used for debugging purposes.  E.g., 'TensorFlow.js Converter v1.0.0'.
+   *
+   * A value of `null` means the model artifacts are generated without any
+   * conversion process (e.g., saved directly from a TensorFlow.js
+   * `LayersModel` instance.)
+   */
+  convertedBy?: string|null;
 }
 
+// TODO(cais): Create interface spec for JSON model artifact, e.g., the format
+// for the content in a model.json file saved by TensorFlow.js. It should be the
+// same as the `ModelArtifacts` interface, but without the binary `weightData`
+// field. The `weightSpec` field should be replaced with a `weightManifest` one.
+
 /**
  * Type definition for handlers of loading operations.
  */