Support loading models from ModelArtifacts and ModelJSON in loadGraphModelSync (#6847)

Let users pass inputs of type ModelArtifacts or [ModelJSON, ArrayBuffer] to loadGraphModelSync instead of requiring them to construct a synchronous IOHandler. In the case of ModelArtifacts, the model weights are contained in the object. In the case of [ModelJSON, ArrayBuffer], the ArrayBuffer is a list of concatenated weights for the model.

This change turns this code:

const modelJson = JSON.parse(modelString);
const weights = ...
modelJson.weightData = weights;
modelJson.weightSpecs = modelJson.weightsManifest[0].weights;
const ioHandler = tf.io.fromMemorySync(modelJson);
const model = tf.loadGraphModelSync(ioHandler);

... into this simpler code

const modelJson = JSON.parse(modelString);
const weights = ...
const model = tf.loadGraphModelSync([modelJson, weights]);

Author: mattsoulanille

tfjs-converter/src/executor/graph_model.ts

@@ -499,23 +499,57 @@ export async function loadGraphModel(
 /**
  * Load a graph model given a synchronous IO handler with a 'load' method.
  *
- * @param modelSource The `io.IOHandlerSync` that loads the model.
+ * @param modelSource The `io.IOHandlerSync` that loads the model, or the
+ *     `io.ModelArtifacts` that encode the model, or a tuple of
+ *     `[io.ModelJSON, ArrayBuffer]` of which the first element encodes the
+ *      model and the second contains the weights.
  *
  * @doc {heading: 'Models', subheading: 'Loading'}
  */
+export function loadGraphModelSync(modelSource: io.IOHandlerSync
+  | io.ModelArtifacts | [io.ModelJSON, /* Weights */ ArrayBuffer]):
+  GraphModel<io.IOHandlerSync> {
 
-export function loadGraphModelSync(modelSource: io.IOHandlerSync):
-    GraphModel<io.IOHandlerSync> {
   if (modelSource == null) {
     throw new Error(
-        'modelUrl in loadGraphModelSync() cannot be null. Please provide a ' +
-        'url or an IOHandler that loads the model');
+        'modelUrl in loadGraphModelSync() cannot be null. Please provide ' +
+        'model artifacts or an IOHandler that loads the model');
   }
-  if (!modelSource.load) {
-    throw new Error(`modelUrl IO Handler ${modelSource} has no load function`);
+
+  let ioHandler: io.IOHandlerSync;
+  if (modelSource instanceof Array) {
+    const [modelJSON, weights] = modelSource;
+    if (!modelJSON) {
+      throw new Error('modelJSON must be the first element of the array');
+    }
+    if (!weights || !(weights instanceof ArrayBuffer)) {
+      throw new Error('An ArrayBuffer of weights must be the second element of'
+                      + ' the array');
+    }
+    if (!('modelTopology' in modelJSON)) {
+      throw new Error('Model JSON is missing \'modelTopology\'');
+    }
+    if (!('weightsManifest' in modelJSON)) {
+      throw new Error('Model JSON is missing \'weightsManifest\'');
+    }
+
+    const weightSpecs = io.getWeightSpecs(modelJSON.weightsManifest);
+    const modelArtifacts = io.getModelArtifactsForJSONSync(modelJSON,
+                                                           weightSpecs,
+                                                           weights);
+    ioHandler = io.fromMemorySync(modelArtifacts);
+  } else if ('load' in modelSource) {
+    // Then modelSource is already an IOHandlerSync.
+    ioHandler = modelSource;
+  } else if ('modelTopology' in modelSource && 'weightSpecs' in modelSource
+      && 'weightData' in modelSource) {
+    // modelSource is of type ModelArtifacts.
+    ioHandler = io.fromMemorySync(modelSource);
+  } else {
+    throw new Error('Unknown model format');
   }
-  const model = new GraphModel(modelSource);
 
+  const model = new GraphModel(ioHandler);
   model.load();
   return model;
 }

tfjs-converter/src/executor/graph_model_test.ts

@@ -425,6 +425,13 @@ describe('loadGraphModel', () => {
 });
 
 describe('loadGraphModelSync', () => {
+  function checkModel(model: GraphModel<unknown>) {
+    expect(model).toBeDefined();
+    const bias = model.weights['Const'][0];
+    expect(bias.dtype).toBe('int32');
+    expect(bias.dataSync()).toEqual(new Int32Array([5]));
+  }
+
   it('Pass a custom io handler', () => {
     const customLoader: tfc.io.IOHandlerSync = {
       load: () => {
@@ -436,10 +443,69 @@ describe('loadGraphModelSync', () => {
       }
     };
     const model = loadGraphModelSync(customLoader);
-    expect(model).toBeDefined();
-    const bias = model.weights['Const'][0];
-    expect(bias.dtype).toBe('int32');
-    expect(bias.dataSync()).toEqual(new Int32Array([5]));
+    checkModel(model);
+  });
+
+  it('Pass the model artifacts directly', () => {
+    const modelArtifacts: tfc.io.ModelArtifacts = {
+      modelTopology: SIMPLE_MODEL,
+      weightSpecs: weightsManifest,
+      weightData: new Int32Array([5]).buffer,
+    };
+    const model = loadGraphModelSync(modelArtifacts);
+    checkModel(model);
+  });
+
+  it('Pass the model JSON and weights', () => {
+    const modelJson: tfc.io.ModelJSON = {
+      modelTopology: SIMPLE_MODEL,
+      weightsManifest: [{paths: [], weights: weightsManifest}],
+    };
+    const weights = new Int32Array([5]).buffer;
+    const model = loadGraphModelSync([modelJson, weights]);
+    checkModel(model);
+  });
+
+  it('Throws an error if ModelJSON is passed without weights', () => {
+    const modelJson: tfc.io.ModelJSON = {
+      modelTopology: SIMPLE_MODEL,
+      weightsManifest: [{paths: [], weights: weightsManifest}],
+    };
+    expect(() => {
+      return loadGraphModelSync([modelJson] as unknown as [io.ModelJSON,
+                                                           ArrayBuffer]);
+    }).toThrowMatching(err =>
+      err.message.includes('weights must be the second element'));
+  });
+
+  it('Throws an error if modelJSON is missing \'modelTopology\'', () => {
+    const badInput = {
+      weightsManifest: [{paths: [] as string[], weights: weightsManifest}],
+    };
+    const weights = new Int32Array([5]).buffer;
+    expect(() => {
+      return loadGraphModelSync([badInput as io.ModelJSON, weights]);
+    }).toThrowMatching(err =>
+      err.message.includes('missing \'modelTopology\''));
+  });
+
+  it('Throws an error if modelJSON is missing \'weightsManifest\'', () => {
+    const badInput = {
+      modelTopology: SIMPLE_MODEL,
+    };
+    const weights = new Int32Array([5]).buffer;
+    expect(() => {
+      return loadGraphModelSync([badInput as io.ModelJSON, weights]);
+    }).toThrowMatching(err =>
+      err.message.includes('missing \'weightsManifest\''));
+  });
+
+  it('Throws an error if modelSource is an unknown format', () => {
+    const badInput = {foo: 'bar'};
+    expect(() => {
+      return loadGraphModelSync(badInput as io.ModelArtifacts);
+    }).toThrowMatching(err =>
+      err.message.includes('Unknown model format'));
   });
 
   it('Expect an error when moderUrl is null', () => {

tfjs-core/src/io/http.ts

@@ -24,7 +24,7 @@
 import {env} from '../environment';
 
 import {assert} from '../util';
-import {concatenateArrayBuffers, getModelArtifactsForJSON, getModelArtifactsInfoForJSON, getModelJSONForModelArtifacts} from './io_utils';
+import {concatenateArrayBuffers, getModelArtifactsForJSON, getModelArtifactsInfoForJSON, getModelJSONForModelArtifacts, getWeightSpecs} from './io_utils';
 import {IORouter, IORouterRegistry} from './router_registry';
 import {IOHandler, LoadOptions, ModelArtifacts, ModelJSON, OnProgressCallback, SaveResult, WeightsManifestConfig, WeightsManifestEntry} from './types';
 import {loadWeightsAsArrayBuffer} from './weights_loader';
@@ -187,10 +187,7 @@ export class HTTPRequest implements IOHandler {
     const [prefix, suffix] = parseUrl(weightPath);
     const pathPrefix = this.weightPathPrefix || prefix;
 
-    const weightSpecs = [];
-    for (const entry of weightsManifest) {
-      weightSpecs.push(...entry.weights);
-    }
+    const weightSpecs = getWeightSpecs(weightsManifest);
 
     const fetchURLs: string[] = [];
     const urlPromises: Array<Promise<string>> = [];

tfjs-core/src/io/io.ts

@@ -22,7 +22,7 @@ import './local_storage';
 
 import {browserFiles} from './browser_files';
 import {browserHTTPRequest, http, isHTTPScheme} from './http';
-import {concatenateArrayBuffers, decodeWeights, encodeWeights, getModelArtifactsForJSON, getModelArtifactsInfoForJSON} from './io_utils';
+import {concatenateArrayBuffers, decodeWeights, encodeWeights, getModelArtifactsForJSON, getModelArtifactsForJSONSync, getModelArtifactsInfoForJSON, getWeightSpecs} from './io_utils';
 import {fromMemory, fromMemorySync, withSaveHandler, withSaveHandlerSync} from './passthrough';
 import {getLoadHandlers, getSaveHandlers, registerLoadRouter, registerSaveRouter} from './router_registry';
 import {IOHandler, IOHandlerSync, LoadHandler, LoadOptions, ModelArtifacts, ModelArtifactsInfo, ModelJSON, ModelStoreManager, OnProgressCallback, RequestDetails, SaveConfig, SaveHandler, SaveResult, TrainingConfig, WeightGroup, WeightsManifestConfig, WeightsManifestEntry} from './types';
@@ -39,8 +39,10 @@ export {
   fromMemorySync,
   getLoadHandlers,
   getModelArtifactsForJSON,
+  getModelArtifactsForJSONSync,
   getModelArtifactsInfoForJSON,
   getSaveHandlers,
+  getWeightSpecs,
   http,
   IOHandler,
   IOHandlerSync,

tfjs-core/src/io/io_utils.ts

@@ -403,19 +403,20 @@ export function getModelJSONForModelArtifacts(
 }
 
 /**
- * Create `ModelArtifacts` from a JSON file.
+ * Create `ModelArtifacts` from a JSON file and weights.
  *
  * @param modelJSON Object containing the parsed JSON of `model.json`
- * @param loadWeights Function that takes the JSON file's weights manifest,
- *     reads weights from the listed path(s), and returns a Promise of the
- *     weight manifest entries along with the weights data.
+ * @param weightSpecs The list of WeightsManifestEntry for the model. Must be
+ *     passed if the modelJSON has a weightsManifest.
+ * @param weightData An ArrayBuffer of weight data for the model corresponding
+ *     to the weights in weightSpecs. Must be passed if the modelJSON has a
+ *     weightsManifest.
  * @returns A Promise of the `ModelArtifacts`, as described by the JSON file.
  */
-export async function getModelArtifactsForJSON(
-    modelJSON: ModelJSON,
-    loadWeights: (weightsManifest: WeightsManifestConfig) => Promise<[
-      /* weightSpecs */ WeightsManifestEntry[], /* weightData */ ArrayBuffer
-    ]>): Promise<ModelArtifacts> {
+export function getModelArtifactsForJSONSync(
+    modelJSON: ModelJSON, weightSpecs?: WeightsManifestEntry[],
+    weightData?: ArrayBuffer): ModelArtifacts {
+
   const modelArtifacts: ModelArtifacts = {
     modelTopology: modelJSON.modelTopology,
     format: modelJSON.format,
@@ -427,8 +428,12 @@ export async function getModelArtifactsForJSON(
     modelArtifacts.trainingConfig = modelJSON.trainingConfig;
   }
   if (modelJSON.weightsManifest != null) {
-    const [weightSpecs, weightData] =
-        await loadWeights(modelJSON.weightsManifest);
+    if (!weightSpecs) {
+      throw new Error('modelJSON has weightsManifest but weightSpecs is null');
+    }
+    if (!weightData) {
+      throw new Error('modelJSON has weightsManifest but weightData is null');
+    }
     modelArtifacts.weightSpecs = weightSpecs;
     modelArtifacts.weightData = weightData;
   }
@@ -445,6 +450,30 @@ export async function getModelArtifactsForJSON(
   return modelArtifacts;
 }
 
+/**
+ * Create `ModelArtifacts` from a JSON file.
+ *
+ * @param modelJSON Object containing the parsed JSON of `model.json`
+ * @param loadWeights Function that takes the JSON file's weights manifest,
+ *     reads weights from the listed path(s), and returns a Promise of the
+ *     weight manifest entries along with the weights data.
+ * @returns A Promise of the `ModelArtifacts`, as described by the JSON file.
+ */
+export async function getModelArtifactsForJSON(
+    modelJSON: ModelJSON,
+    loadWeights: (weightsManifest: WeightsManifestConfig) => Promise<[
+      /* weightSpecs */ WeightsManifestEntry[], /* weightData */ ArrayBuffer
+    ]>): Promise<ModelArtifacts> {
+  let weightSpecs: WeightsManifestEntry[] | undefined;
+  let weightData: ArrayBuffer | undefined;
+
+  if (modelJSON.weightsManifest != null) {
+    [weightSpecs, weightData] = await loadWeights(modelJSON.weightsManifest);
+  }
+
+  return getModelArtifactsForJSONSync(modelJSON, weightSpecs, weightData);
+}
+
 /**
  * Populate ModelArtifactsInfo fields for a model with JSON topology.
  * @param modelArtifacts
@@ -471,6 +500,22 @@ export function getModelArtifactsInfoForJSON(modelArtifacts: ModelArtifacts):
   };
 }
 
+/**
+ * Concatenate the weights stored in a WeightsManifestConfig into a list of
+ * WeightsManifestEntry
+ *
+ * @param weightsManifest The WeightsManifestConfig to extract weights from.
+ * @returns A list of WeightsManifestEntry of the weights in the weightsManifest
+ */
+export function getWeightSpecs(weightsManifest: WeightsManifestConfig):
+    WeightsManifestEntry[] {
+  const weightSpecs: WeightsManifestEntry[] = [];
+  for (const entry of weightsManifest) {
+    weightSpecs.push(...entry.weights);
+  }
+  return weightSpecs;
+}
+
 /**
  * Computes mantisa table for casting Float16 to Float32
  * See http://www.fox-toolkit.org/ftp/fasthalffloatconversion.pdf