diff --git a/tfjs-core/src/base.ts b/tfjs-core/src/base.ts index a038a7cfe6a..a36679efc7a 100644 --- a/tfjs-core/src/base.ts +++ b/tfjs-core/src/base.ts @@ -41,7 +41,7 @@ import * as test_util from './test_util'; import * as util from './util'; import {version} from './version'; -export {InferenceModel, MetaGraph, MetaGraphInfo, ModelPredictConfig, ModelTensorInfo, SavedModelTensorInfo, SignatureDef, SignatureDefInfo} from './model_types'; +export {InferenceModel, MetaGraph, MetaGraphInfo, ModelPredictConfig, ModelTensorInfo, SavedModelTensorInfo, SignatureDef, SignatureDefEntry, SignatureDefInfo} from './model_types'; // Optimizers. export {AdadeltaOptimizer} from './optimizers/adadelta_optimizer'; export {AdagradOptimizer} from './optimizers/adagrad_optimizer'; diff --git a/tfjs-core/src/model_types.ts b/tfjs-core/src/model_types.ts index 3c388a96b84..9484e978a0b 100644 --- a/tfjs-core/src/model_types.ts +++ b/tfjs-core/src/model_types.ts @@ -41,6 +41,8 @@ export interface ModelTensorInfo { shape?: number[]; // Data type of the tensor. dtype: DataType; + // TensorFlow native Data type of the tensor. + tfDtype?: string; } /** @@ -138,12 +140,17 @@ export interface MetaGraph { signatureDefs: SignatureDef; } +/** + * Interface for SavedModel/GraphModel SignatureDef entry. + */ +export interface SignatureDefEntry { + inputs: {[key: string]: ModelTensorInfo}; + outputs: {[key: string]: ModelTensorInfo}; +} + /** * Interface for SavedModel/GraphModel SignatureDef info. */ export interface SignatureDef { - [key: string]: { - inputs: {[key: string]: ModelTensorInfo}; - outputs: {[key: string]: ModelTensorInfo}; - }; + [key: string]: SignatureDefEntry; } diff --git a/tfjs-node-gpu/test_objects/saved_model/uint8_multiply/saved_model.pb b/tfjs-node-gpu/test_objects/saved_model/uint8_multiply/saved_model.pb new file mode 100644 index 00000000000..19587a3fdc8 Binary files /dev/null and b/tfjs-node-gpu/test_objects/saved_model/uint8_multiply/saved_model.pb differ diff --git a/tfjs-node-gpu/test_objects/saved_model/uint8_multiply/variables/variables.data-00000-of-00001 b/tfjs-node-gpu/test_objects/saved_model/uint8_multiply/variables/variables.data-00000-of-00001 new file mode 100644 index 00000000000..c42273995ff Binary files /dev/null and b/tfjs-node-gpu/test_objects/saved_model/uint8_multiply/variables/variables.data-00000-of-00001 differ diff --git a/tfjs-node-gpu/test_objects/saved_model/uint8_multiply/variables/variables.index b/tfjs-node-gpu/test_objects/saved_model/uint8_multiply/variables/variables.index new file mode 100644 index 00000000000..8ad1e6e034b Binary files /dev/null and b/tfjs-node-gpu/test_objects/saved_model/uint8_multiply/variables/variables.index differ diff --git a/tfjs-node/python/unint8_model.py b/tfjs-node/python/unint8_model.py new file mode 100644 index 00000000000..733e68fc259 --- /dev/null +++ b/tfjs-node/python/unint8_model.py @@ -0,0 +1,35 @@ + # Copyright 2020 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Python script for creating Tensorflow SavedModel with UINT8 input.""" + +import os + +import tensorflow as tf +from tensorflow.python.eager import def_function +from tensorflow.python.framework import constant_op +from tensorflow.python.ops import variables +from tensorflow.python.training.tracking import tracking +from tensorflow.python.saved_model.save import save + +"""Test a basic model with functions to make sure functions are inlined.""" +input_data = constant_op.constant(1, shape=[1], dtype=tf.uint8) +root = tracking.AutoTrackable() +root.v1 = variables.Variable(3) +root.v2 = variables.Variable(2) +root.f = def_function.function(lambda x: root.v1 * root.v2 * tf.cast(x, tf.int32)) +to_save = root.f.get_concrete_function(input_data) + +save_dir = os.path.join('..', 'test_objects', 'saved_model', 'uint8_multiply') +save(root, save_dir, to_save) diff --git a/tfjs-node/src/nodejs_kernel_backend.ts b/tfjs-node/src/nodejs_kernel_backend.ts index 466f68e30b8..b1f09694cdc 100644 --- a/tfjs-node/src/nodejs_kernel_backend.ts +++ b/tfjs-node/src/nodejs_kernel_backend.ts @@ -16,7 +16,7 @@ */ import * as tf from '@tensorflow/tfjs'; -import {backend_util, BackendTimingInfo, DataId, DataType, fill, KernelBackend, ones, Rank, rsqrt, Scalar, scalar, ScalarLike, ShapeMap, Tensor, Tensor1D, tensor1d, Tensor2D, tensor2d, Tensor3D, Tensor4D, Tensor5D, TensorInfo, tidy, util} from '@tensorflow/tfjs'; +import {backend_util, BackendTimingInfo, DataId, DataType, fill, KernelBackend, ModelTensorInfo, ones, Rank, rsqrt, Scalar, scalar, ScalarLike, ShapeMap, Tensor, Tensor1D, tensor1d, Tensor2D, tensor2d, Tensor3D, Tensor4D, Tensor5D, TensorInfo, tidy, util} from '@tensorflow/tfjs'; import {isArray, isNullOrUndefined} from 'util'; import {Int64Scalar} from './int64_tensors'; @@ -1903,11 +1903,27 @@ export class NodeJSKernelBackend extends KernelBackend { return this.binding.loadSavedModel(path, tags); } + private getMappedInputTensorIds( + inputs: Tensor[], inputTensorInfos: ModelTensorInfo[]) { + const tensorIds = this.getInputTensorIds(inputs); + for (let i = 0; i < inputs.length; i++) { + if (inputTensorInfos[i] != null && + inputTensorInfos[i].tfDtype === 'DT_UINT8') { + const data = Uint8Array.from(inputs[i].dataSync()); + const inputTensorId = this.binding.createTensor( + inputs[i].shape, this.binding.TF_UINT8, data); + tensorIds[i] = inputTensorId; + } + } + return tensorIds; + } + runSavedModel( - id: number, inputs: Tensor[], inputOpNames: string[], + id: number, inputs: Tensor[], inputTensorInfos: ModelTensorInfo[], outputOpNames: string[]): Tensor[] { const outputMetadata = this.binding.runSavedModel( - id, this.getInputTensorIds(inputs), inputOpNames.join(','), + id, this.getMappedInputTensorIds(inputs, inputTensorInfos), + inputTensorInfos.map(info => info.name).join(','), outputOpNames.join(',')); return outputMetadata.map(m => this.createOutputTensor(m)); } diff --git a/tfjs-node/src/saved_model.ts b/tfjs-node/src/saved_model.ts index d9ac5738279..81f1f191357 100644 --- a/tfjs-node/src/saved_model.ts +++ b/tfjs-node/src/saved_model.ts @@ -15,9 +15,10 @@ * ============================================================================= */ -import {DataType, InferenceModel, MetaGraph, ModelPredictConfig, ModelTensorInfo, NamedTensorMap, SignatureDef, Tensor, util} from '@tensorflow/tfjs'; +import {DataType, InferenceModel, MetaGraph, ModelPredictConfig, ModelTensorInfo, NamedTensorMap, SignatureDef, SignatureDefEntry, Tensor, util} from '@tensorflow/tfjs'; import * as fs from 'fs'; import {promisify} from 'util'; + import {ensureTensorflowBackend, nodeBackend, NodeJSKernelBackend} from './nodejs_kernel_backend'; const readFile = promisify(fs.readFile); @@ -88,8 +89,8 @@ export async function getMetaGraphsFromSavedModel(path: string): // Get SavedModel proto message const modelMessage = await readSavedModelProto(path); - // A SavedModel might have multiple MetaGraphs, identified by tags. Each - // MetaGraph also has it's own signatureDefs. + // A SavedModel might have multiple MetaGraphs, identified by tags. + // Each MetaGraph also has it's own signatureDefs. const metaGraphList = modelMessage.getMetaGraphsList(); for (let i = 0; i < metaGraphList.length; i++) { const metaGraph = {} as MetaGraph; @@ -124,9 +125,10 @@ export async function getMetaGraphsFromSavedModel(path: string): } const inputTensor = inputsMapMessage.get(inputsMapKey.value); const inputTensorInfo = {} as ModelTensorInfo; - inputTensorInfo.dtype = mapTFDtypeToJSDtype( - getEnumKeyFromValue(messages.DataType, inputTensor.getDtype())); - + const dtype = + getEnumKeyFromValue(messages.DataType, inputTensor.getDtype()); + inputTensorInfo.dtype = mapTFDtypeToJSDtype(dtype); + inputTensorInfo.tfDtype = dtype; inputTensorInfo.name = inputTensor.getName(); inputTensorInfo.shape = inputTensor.getTensorShape().getDimList(); inputs[inputsMapKey.value] = inputTensorInfo; @@ -143,8 +145,10 @@ export async function getMetaGraphsFromSavedModel(path: string): } const outputTensor = outputsMapMessage.get(outputsMapKey.value); const outputTensorInfo = {} as ModelTensorInfo; - outputTensorInfo.dtype = mapTFDtypeToJSDtype( - getEnumKeyFromValue(messages.DataType, outputTensor.getDtype())); + const dtype = + getEnumKeyFromValue(messages.DataType, outputTensor.getDtype()); + outputTensorInfo.dtype = mapTFDtypeToJSDtype(dtype); + outputTensorInfo.tfDtype = dtype; outputTensorInfo.name = outputTensor.getName(); outputTensorInfo.shape = outputTensor.getTensorShape().getDimList(); outputs[outputsMapKey.value] = outputTensorInfo; @@ -160,39 +164,24 @@ export async function getMetaGraphsFromSavedModel(path: string): } /** - * Get input and output node names from SavedModel metagraphs info. The - * input.output node names will be used when executing a SavedModel signature. + * Get SignatureDefEntry from SavedModel metagraphs info. The SignatureDefEntry + * will be used when executing a SavedModel signature. * * @param savedModelInfo The MetaGraphInfo array loaded through * getMetaGraphsFromSavedModel(). * @param tags The tags of the MetaGraph to get input/output node names from. * @param signature The signature to get input/output node names from. */ -export function getInputAndOutputNodeNameFromMetaGraphInfo( - savedModelInfo: MetaGraph[], tags: string[], signature: string) { +export function getSignatureDefEntryFromMetaGraphInfo( + savedModelInfo: MetaGraph[], tags: string[], + signature: string): SignatureDefEntry { for (let i = 0; i < savedModelInfo.length; i++) { const metaGraphInfo = savedModelInfo[i]; if (stringArraysHaveSameElements(tags, metaGraphInfo.tags)) { if (metaGraphInfo.signatureDefs[signature] == null) { throw new Error('The SavedModel does not have signature: ' + signature); } - const inputNodeNames: {[key: string]: string} = {}; - const outputNodeNames: {[key: string]: string} = {}; - for (const signatureDef of Object.keys(metaGraphInfo.signatureDefs)) { - if (signatureDef === signature) { - for (const tensorName of Object.keys( - metaGraphInfo.signatureDefs[signature].inputs)) { - inputNodeNames[tensorName] = - metaGraphInfo.signatureDefs[signature].inputs[tensorName].name; - } - for (const tensorName of Object.keys( - metaGraphInfo.signatureDefs[signature].outputs)) { - outputNodeNames[tensorName] = - metaGraphInfo.signatureDefs[signature].outputs[tensorName].name; - } - } - } - return [inputNodeNames, outputNodeNames]; + return metaGraphInfo.signatureDefs[signature]; } } throw new Error(`The SavedModel does not have tags: ${tags}`); @@ -206,11 +195,10 @@ export function getInputAndOutputNodeNameFromMetaGraphInfo( */ export class TFSavedModel implements InferenceModel { private disposed = false; - + private outputNodeNames_: {[key: string]: string}; constructor( private sessionId: number, private jsid: number, - private inputNodeNames: {[key: string]: string}, - private outputNodeNames: {[key: string]: string}, + private signature: SignatureDefEntry, private backend: NodeJSKernelBackend) {} /** @@ -254,6 +242,19 @@ export class TFSavedModel implements InferenceModel { } } + get outputNodeNames() { + if (this.outputNodeNames_ != null) { + return this.outputNodeNames_; + } + this.outputNodeNames_ = + Object.keys(this.signature.outputs) + .reduce((names: {[key: string]: string}, key: string) => { + names[key] = this.signature.outputs[key].name; + return names; + }, {}); + return this.outputNodeNames_; + } + /** * Execute the inference for the input tensors. * @@ -287,16 +288,16 @@ export class TFSavedModel implements InferenceModel { if (inputs instanceof Tensor) { inputTensors.push(inputs); const result = this.backend.runSavedModel( - this.sessionId, inputTensors, Object.values(this.inputNodeNames), + this.sessionId, inputTensors, Object.values(this.signature.inputs), Object.values(this.outputNodeNames)); return result.length > 1 ? result : result[0]; } else if (Array.isArray(inputs)) { inputTensors = inputs; return this.backend.runSavedModel( - this.sessionId, inputTensors, Object.values(this.inputNodeNames), + this.sessionId, inputTensors, Object.values(this.signature.inputs), Object.values(this.outputNodeNames)); } else { - const inputTensorNames = Object.keys(this.inputNodeNames); + const inputTensorNames = Object.keys(this.signature.inputs); const providedInputNames = Object.keys(inputs); if (!stringArraysHaveSameElements( inputTensorNames, providedInputNames)) { @@ -304,10 +305,10 @@ export class TFSavedModel implements InferenceModel { inputTensorNames.join()}, however the provided input names are ${ providedInputNames.join()}.`); } - const inputNodeNamesArray = []; + const inputNodeNamesArray: ModelTensorInfo[] = []; for (let i = 0; i < inputTensorNames.length; i++) { inputTensors.push(inputs[inputTensorNames[i]]); - inputNodeNamesArray.push(this.inputNodeNames[inputTensorNames[i]]); + inputNodeNamesArray.push(this.signature.inputs[inputTensorNames[i]]); } const outputTensorNames = Object.keys(this.outputNodeNames); const outputNodeNamesArray = []; @@ -387,9 +388,8 @@ export async function loadSavedModel( const backend = nodeBackend(); const savedModelInfo = await getMetaGraphsFromSavedModel(path); - const [inputNodeNames, outputNodeNames] = - getInputAndOutputNodeNameFromMetaGraphInfo( - savedModelInfo, tags, signature); + const signatureDefEntry = + getSignatureDefEntryFromMetaGraphInfo(savedModelInfo, tags, signature); let sessionId: number; @@ -407,7 +407,7 @@ export async function loadSavedModel( } const id = nextTFSavedModelId++; const savedModel = - new TFSavedModel(sessionId, id, inputNodeNames, outputNodeNames, backend); + new TFSavedModel(sessionId, id, signatureDefEntry, backend); loadedSavedModelPathMap.set(id, {path, tags, sessionId}); return savedModel; } @@ -431,6 +431,7 @@ function mapTFDtypeToJSDtype(tfDtype: string): DataType { case 'DT_FLOAT': return 'float32'; case 'DT_INT32': + case 'DT_UINT8': return 'int32'; case 'DT_BOOL': return 'bool'; diff --git a/tfjs-node/src/saved_model_test.ts b/tfjs-node/src/saved_model_test.ts index 25a9cbf3489..3815020d0f1 100644 --- a/tfjs-node/src/saved_model_test.ts +++ b/tfjs-node/src/saved_model_test.ts @@ -18,7 +18,7 @@ import {NamedTensorMap, test_util} from '@tensorflow/tfjs'; import * as tf from './index'; import {nodeBackend} from './nodejs_kernel_backend'; -import {getEnumKeyFromValue, getInputAndOutputNodeNameFromMetaGraphInfo, readSavedModelProto} from './saved_model'; +import {getEnumKeyFromValue, getSignatureDefEntryFromMetaGraphInfo, readSavedModelProto} from './saved_model'; // tslint:disable-next-line:no-require-imports const messages = require('./proto/api_pb'); @@ -187,11 +187,11 @@ describe('SavedModel', () => { it('get input and output node names from SavedModel metagraphs', async () => { const modelInfo = await tf.node.getMetaGraphsFromSavedModel( './test_objects/saved_model/times_three_float'); - const inputAndOutputNodeNames = getInputAndOutputNodeNameFromMetaGraphInfo( + const signature = getSignatureDefEntryFromMetaGraphInfo( modelInfo, ['serve'], 'serving_default'); - expect(inputAndOutputNodeNames.length).toBe(2); - expect(inputAndOutputNodeNames[0]['x']).toBe('serving_default_x:0'); - expect(inputAndOutputNodeNames[1]['output_0']) + expect(Object.keys(signature).length).toBe(2); + expect(signature.inputs['x'].name).toBe('serving_default_x:0'); + expect(signature.outputs['output_0'].name) .toBe('StatefulPartitionedCall:0'); }); @@ -394,6 +394,17 @@ describe('SavedModel', () => { done(); } }); + it('execute model with uint8 input', async () => { + const model = await tf.node.loadSavedModel( + './test_objects/saved_model/uint8_multiply', ['serve'], + 'serving_default'); + const input = tf.scalar(3, 'int32'); + const output = model.predict(input) as tf.Tensor; + expect(output.shape).toEqual([]); + expect(output.dtype).toBe('int32'); + test_util.expectArraysClose(await output.data(), [18]); + model.dispose(); + }); it('execute model int times two', async () => { const model = await tf.node.loadSavedModel( diff --git a/tfjs-node/test_objects/saved_model/uint8_multiply/saved_model.pb b/tfjs-node/test_objects/saved_model/uint8_multiply/saved_model.pb new file mode 100644 index 00000000000..19587a3fdc8 Binary files /dev/null and b/tfjs-node/test_objects/saved_model/uint8_multiply/saved_model.pb differ diff --git a/tfjs-node/test_objects/saved_model/uint8_multiply/variables/variables.data-00000-of-00001 b/tfjs-node/test_objects/saved_model/uint8_multiply/variables/variables.data-00000-of-00001 new file mode 100644 index 00000000000..c42273995ff Binary files /dev/null and b/tfjs-node/test_objects/saved_model/uint8_multiply/variables/variables.data-00000-of-00001 differ diff --git a/tfjs-node/test_objects/saved_model/uint8_multiply/variables/variables.index b/tfjs-node/test_objects/saved_model/uint8_multiply/variables/variables.index new file mode 100644 index 00000000000..8ad1e6e034b Binary files /dev/null and b/tfjs-node/test_objects/saved_model/uint8_multiply/variables/variables.index differ