Provide types to enforce the structure of model.json files. (tensorflow#1597)

We were reminded of this by discussion on tensorflow#1596 (and, related, tensorflow/tfjs-converter#320).

My view is that those two PRs should make it into 1.0.  This one is less urgent because it's a non-breaking change that just helps enforce our intentions.

Obviously, this PR can't enforce anything on the Python side.  Ultimately we should probably do all this via a language-independent JSON Schema.

INTERNAL

Author: davidsoergel

src/io/browser_files.ts

@@ -23,7 +23,7 @@
 import {ENV} from '../environment';
 import {basename, concatenateArrayBuffers, getModelArtifactsInfoForJSON} from './io_utils';
 import {IORouter, IORouterRegistry} from './router_registry';
-import {IOHandler, ModelArtifacts, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
+import {IOHandler, ModelArtifacts, ModelJSON, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
 
 const DEFAULT_FILE_NAME_PREFIX = 'model';
 const DEFAULT_JSON_EXTENSION_NAME = '.json';
@@ -71,7 +71,7 @@ export class BrowserDownloads implements IOHandler {
         paths: ['./' + this.weightDataFileName],
         weights: modelArtifacts.weightSpecs
       }];
-      const modelTopologyAndWeightManifest = {
+      const modelTopologyAndWeightManifest: ModelJSON = {
         modelTopology: modelArtifacts.modelTopology,
         format: modelArtifacts.format,
         generatedBy: modelArtifacts.generatedBy,
@@ -127,8 +127,8 @@ class BrowserFiles implements IOHandler {
       const jsonReader = new FileReader();
       jsonReader.onload = (event: Event) => {
         // tslint:disable-next-line:no-any
-        const modelJSON = JSON.parse((event.target as any).result);
-        const modelTopology = modelJSON.modelTopology as {};
+        const modelJSON = JSON.parse((event.target as any).result) as ModelJSON;
+        const modelTopology = modelJSON.modelTopology;
         if (modelTopology == null) {
           reject(new Error(
               `modelTopology field is missing from file ${jsonFile.name}`));
@@ -139,8 +139,7 @@ class BrowserFiles implements IOHandler {
           resolve({modelTopology});
         }
 
-        const weightsManifest =
-            modelJSON.weightsManifest as WeightsManifestConfig;
+        const weightsManifest = modelJSON.weightsManifest;
         if (weightsManifest == null) {
           reject(new Error(
               `weightManifest field is missing from file ${jsonFile.name}`));

src/io/browser_http.ts

@@ -22,10 +22,9 @@
  */
 
 import {assert} from '../util';
-
 import {concatenateArrayBuffers, getModelArtifactsInfoForJSON} from './io_utils';
 import {IORouter, IORouterRegistry} from './router_registry';
-import {IOHandler, LoadOptions, ModelArtifacts, OnProgressCallback, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
+import {IOHandler, LoadOptions, ModelArtifacts, ModelJSON, OnProgressCallback, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
 import {loadWeightsAsArrayBuffer} from './weights_loader';
 
 const OCTET_STREAM_MIME_TYPE = 'application/octet-stream';
@@ -112,7 +111,7 @@ export class BrowserHTTPRequest implements IOHandler {
       paths: ['./model.weights.bin'],
       weights: modelArtifacts.weightSpecs,
     }];
-    const modelTopologyAndWeightManifest = {
+    const modelTopologyAndWeightManifest: ModelJSON = {
       modelTopology: modelArtifacts.modelTopology,
       format: modelArtifacts.format,
       generatedBy: modelArtifacts.generatedBy,
@@ -166,7 +165,7 @@ export class BrowserHTTPRequest implements IOHandler {
           `${modelConfigRequest.status}. Please verify this URL points to ` +
           `the model JSON of the model to load.`);
     }
-    let modelConfig;
+    let modelConfig: ModelJSON;
     try {
       modelConfig = await modelConfigRequest.json();
     } catch (e) {
@@ -186,9 +185,8 @@ export class BrowserHTTPRequest implements IOHandler {
       }
       throw new Error(message);
     }
-
-    const modelTopology = modelConfig['modelTopology'];
-    const weightsManifest = modelConfig['weightsManifest'];
+    const modelTopology = modelConfig.modelTopology;
+    const weightsManifest = modelConfig.weightsManifest;
 
     // We do not allow both modelTopology and weightsManifest to be missing.
     if (modelTopology == null && weightsManifest == null) {
@@ -200,8 +198,6 @@ export class BrowserHTTPRequest implements IOHandler {
     let weightSpecs: WeightsManifestEntry[];
     let weightData: ArrayBuffer;
     if (weightsManifest != null) {
-      const weightsManifest =
-          modelConfig['weightsManifest'] as WeightsManifestConfig;
       const results = await this.loadWeights(weightsManifest);
       [weightSpecs, weightData] = results;
     }

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, ModelFormat, OnProgressCallback, SaveConfig, SaveHandler, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
+import {IOHandler, LoadHandler, LoadOptions, ModelArtifacts, ModelJSON, ModelStoreManager, OnProgressCallback, SaveConfig, SaveHandler, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
 import {loadWeights, weightsLoaderFactory} from './weights_loader';
 
 export {copyModel, listModels, moveModel, removeModel} from './model_management';
@@ -45,7 +45,7 @@ export {
   LoadOptions,
   loadWeights,
   ModelArtifacts,
-  ModelFormat,
+  ModelJSON,
   ModelStoreManager,
   OnProgressCallback,
   registerLoadRouter,

src/io/local_storage.ts

@@ -20,7 +20,7 @@ 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, ModelFormat, SaveResult} from './types';
+import {IOHandler, ModelArtifacts, ModelArtifactsInfo, ModelStoreManager, SaveResult} from './types';
 
 const PATH_SEPARATOR = '/';
 const PATH_PREFIX = 'tensorflowjs_models';
@@ -57,16 +57,20 @@ export function purgeLocalStorageArtifacts(): string[] {
   return purgedModelPaths;
 }
 
-function getModelKeys(path: string):
-    {info: string, topology: string, weightSpecs: string, weightData: string,
-     modelMetadata: string} {
+function getModelKeys(path: 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),
-    modelMetadata: [
-        PATH_PREFIX, path, MODEL_METADATA_SUFFIX].join(PATH_SEPARATOR)
+    modelMetadata:
+        [PATH_PREFIX, path, MODEL_METADATA_SUFFIX].join(PATH_SEPARATOR)
   };
 }
 
@@ -218,12 +222,9 @@ export class BrowserLocalStorage implements IOHandler {
     // 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;
+      const metadata = JSON.parse(metadataString) as
+          {format: string, generatedBy: string, convertedBy: string};
+      out.format = metadata['format'];
       out.generatedBy = metadata['generatedBy'];
       out.convertedBy = metadata['convertedBy'];
     }

src/io/types.ts

@@ -156,22 +156,23 @@ export declare interface ModelArtifactsInfo {
   weightDataBytes?: number;
 }
 
-export declare type ModelFormat = 'graph-model'|'layers-model';
-
 /**
  * The serialized artifacts of a model, including topology and weights.
  *
  * The `modelTopology`, `weightSpecs` and `weightData` fields of this interface
  * are optional, in order to support topology- or weights-only saving and
  * loading.
+ *
+ * Note this interface is used internally in IOHandlers.  For the file format
+ * written to disk as `model.json`, see `ModelJSON`.
  */
 export declare interface ModelArtifacts {
   /**
    * Model topology.
    *
    * For Keras-style `tf.Model`s, this is a JSON object.
-   * For TensorFlow-style models (e.g., `FrozenModel`), this is a binary buffer
-   * carrying the `GraphDef` protocol buffer.
+   * For TensorFlow-style models (e.g., `SavedModel`), this is the JSON
+   * encoding of the `GraphDef` protocol buffer.
    */
   modelTopology?: {}|ArrayBuffer;
 
@@ -192,7 +193,7 @@ export declare interface ModelArtifacts {
    * Hard-coded format name for models saved from TensorFlow.js or converted
    * by TensorFlow.js Converter.
    */
-  format?: ModelFormat;
+  format?: string;
 
   /**
    * What library is responsible for originally generating this artifact.
@@ -209,15 +210,62 @@ export declare interface ModelArtifacts {
    *
    * A value of `null` means the model artifacts are generated without any
    * conversion process (e.g., saved directly from a TensorFlow.js
-   * `LayersModel` instance.)
+   * `tf.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.
+/**
+ * The on-disk format of the `model.json` file.
+ *
+ * TF.js 1.0 always populates the optional fields when writing model.json.
+ * Prior versions did not provide those fields.
+ */
+export declare interface ModelJSON {
+  /**
+   * Model topology.
+   *
+   * For Keras-style `tf.Model`s, this is a JSON object.
+   * For TensorFlow-style models (e.g., `SavedModel`), this is the JSON
+   * encoding of the `GraphDef` protocol buffer.
+   */
+  modelTopology: {};
+
+  /**
+   * Weights manifest.
+   *
+   * The weights manifest consists of an ordered list of weight-manifest
+   * groups. Each weight-manifest group consists of a number of weight values
+   * stored in a number of paths. See the documentation of
+   * `WeightsManifestConfig` for more details.
+   */
+  weightsManifest: WeightsManifestConfig;
+
+  /**
+   * Hard-coded format name for models saved from TensorFlow.js or converted
+   * by TensorFlow.js Converter.
+   */
+  format?: string;
+
+  /**
+   * 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
+   * `tf.LayersModel` instance.)
+   */
+  convertedBy?: string|null;
+}
 
 /**
  * Type definition for handlers of loading operations.