license 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 ) ; const messages = require ( './proto/api_pb' ) ; const SAVED_MODEL_FILE_NAME = '/saved_model.pb' ; const SAVED_MODEL_INIT_OP_KEY = '__saved_model_init_op' ; const loadedSavedModelPathMap = new Map < number , { path : string ; tags : string [ ] ; sessionId : number ; } > ( ) ; let nextTFSavedModelId = 0 ; param object param value export function getEnumKeyFromValue ( object , value ) { return Object . keys ( object ) . find ( key => object [ key ] === value ) ; } param path export async function readSavedModelProto ( path ) { try { fs . accessSync ( path + SAVED_MODEL_FILE_NAME , fs . constants . R_OK ) ; } catch ( error ) { throw new Error ( 'There is no saved_model.pb file in the directory: ' + path ) ; } const modelFile = await readFile ( path + SAVED_MODEL_FILE_NAME ) ; const array = new Uint8Array ( modelFile ) ; return messages . SavedModel . deserializeBinary ( array ) ; } param path doc export async function getMetaGraphsFromSavedModel ( path ) { const result = [ ] ; const modelMessage = await readSavedModelProto ( path ) ; const metaGraphList = modelMessage . getMetaGraphsList ( ) ; for ( let i = 0 ; i < metaGraphList . length ; i ++ ) { const metaGraph = { } as MetaGraph ; const tags = metaGraphList [ i ] . getMetaInfoDef ( ) . getTagsList ( ) ; metaGraph . tags = tags ; const signatureDef = { } ; const signatureDefMap = metaGraphList [ i ] . getSignatureDefMap ( ) ; const signatureDefKeys = signatureDefMap . keys ( ) ; while ( true ) { const key = signatureDefKeys . next ( ) ; if ( key . done ) { break ; } if ( key . value === SAVED_MODEL_INIT_OP_KEY ) { continue ; } const signatureDefEntry = signatureDefMap . get ( key . value ) ; const inputsMapMessage = signatureDefEntry . getInputsMap ( ) ; const inputsMapKeys = inputsMapMessage . keys ( ) ; const inputs = { } ; while ( true ) { const inputsMapKey = inputsMapKeys . next ( ) ; if ( inputsMapKey . done ) { break ; } const inputTensor = inputsMapMessage . get ( inputsMapKey . value ) ; const inputTensorInfo = { } as ModelTensorInfo ; 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 ; } const outputsMapMessage = signatureDefEntry . getOutputsMap ( ) ; const outputsMapKeys = outputsMapMessage . keys ( ) ; const outputs = { } ; while ( true ) { const outputsMapKey = outputsMapKeys . next ( ) ; if ( outputsMapKey . done ) { break ; } const outputTensor = outputsMapMessage . get ( outputsMapKey . value ) ; const outputTensorInfo = { } as ModelTensorInfo ; 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 ; } signatureDef [ key . value ] = { inputs , outputs } ; } metaGraph . signatureDefs = signatureDef ; result . push ( metaGraph ) ; } return result ; } param savedModelInfo param tags param signature export function getSignatureDefEntryFromMetaGraphInfo ( savedModelInfo , tags , signature ) { 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 ) ; } return metaGraphInfo . signatureDefs [ signature ] ; } } throw new Error ( `The SavedModel does not have tags: ${ tags }` ) ; } doc export class TFSavedModel implements InferenceModel { private disposed = false ; private outputNodeNames_ : { [ key ] : string ; } ; constructor ( private sessionId , private jsid , private signature , private backend ) { } doc get inputs ( ) : ModelTensorInfo [ ] { const entries = this . signature . inputs ; const results = Object . keys ( entries ) . map ( ( key ) => entries [ key ] ) ; results . forEach ( ( info ) => { info . name = info . name . replace ( /:0$/ , '' ) ; } ) ; return results ; } doc get outputs ( ) : ModelTensorInfo [ ] { const entries = this . signature . outputs ; const results = Object . keys ( entries ) . map ( ( key ) => entries [ key ] ) ; results . forEach ( ( info ) => { info . name = info . name . replace ( /:0$/ , '' ) ; } ) ; return results ; } doc dispose ( ) { if ( ! this . disposed ) { this . disposed = true ; loadedSavedModelPathMap . delete ( this . jsid ) ; for ( const id of Array . from ( loadedSavedModelPathMap . keys ( ) ) ) { const value = loadedSavedModelPathMap . get ( id ) ; if ( value . sessionId === this . sessionId ) { return ; } } this . backend . deleteSavedModel ( this . sessionId ) ; } else { throw new Error ( 'This SavedModel has already been deleted.' ) ; } } get outputNodeNames ( ) { if ( this . outputNodeNames_ != null ) { return this . outputNodeNames_ ; } this . outputNodeNames_ = Object . keys ( this . signature . outputs ) . reduce ( ( names , key ) => { names [ key ] = this . signature . outputs [ key ] . name ; return names ; } , { } ) ; return this . outputNodeNames_ ; } param input param config returns doc predict ( inputs , config ? ) { if ( this . disposed ) { throw new Error ( 'The TFSavedModel has already been deleted!' ) ; } else { let inputTensors = [ ] ; if ( inputs instanceof Tensor ) { inputTensors . push ( inputs ) ; const result = this . backend . runSavedModel ( 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 . signature . inputs ) , Object . values ( this . outputNodeNames ) ) ; } else { const inputTensorNames = Object . keys ( this . signature . inputs ) ; const providedInputNames = Object . keys ( inputs ) ; if ( ! stringArraysHaveSameElements ( inputTensorNames , providedInputNames ) ) { throw new Error ( `The model signatureDef input names are ${ inputTensorNames . join ( ) }, however the provided input names are ${ providedInputNames . join ( ) }.` ) ; } const inputNodeNamesArray = [ ] ; for ( let i = 0 ; i < inputTensorNames . length ; i ++ ) { inputTensors . push ( inputs [ inputTensorNames [ i ] ] ) ; inputNodeNamesArray . push ( this . signature . inputs [ inputTensorNames [ i ] ] ) ; } const outputTensorNames = Object . keys ( this . outputNodeNames ) ; const outputNodeNamesArray = [ ] ; for ( let i = 0 ; i < outputTensorNames . length ; i ++ ) { outputNodeNamesArray . push ( this . outputNodeNames [ outputTensorNames [ i ] ] ) ; } const outputTensors = this . backend . runSavedModel ( this . sessionId , inputTensors , inputNodeNamesArray , outputNodeNamesArray ) ; util . assert ( outputTensors . length === outputNodeNamesArray . length , ( ) => 'Output tensors do not match output node names, ' + `receive ${ outputTensors . length }) output tensors but ` + `there are ${ this . outputNodeNames . length } output nodes.` ) ; const outputMap = { } ; for ( let i = 0 ; i < outputTensorNames . length ; i ++ ) { outputMap [ outputTensorNames [ i ] ] = outputTensors [ i ] ; } return outputMap ; } } } param input param outputs returns doc execute ( inputs , outputs ) { throw new Error ( 'execute() of TFSavedModel is not supported yet.' ) ; } } param path param tags param signature doc export async function loadSavedModel ( path , tags = [ 'serve' ] , signature = 'serving_default' ) { ensureTensorflowBackend ( ) ; const backend = nodeBackend ( ) ; const savedModelInfo = await getMetaGraphsFromSavedModel ( path ) ; const signatureDefEntry = getSignatureDefEntryFromMetaGraphInfo ( savedModelInfo , tags , signature ) ; let sessionId ; for ( const id of Array . from ( loadedSavedModelPathMap . keys ( ) ) ) { const modelInfo = loadedSavedModelPathMap . get ( id ) ; if ( modelInfo . path === path && stringArraysHaveSameElements ( modelInfo . tags , tags ) ) { sessionId = modelInfo . sessionId ; } } if ( sessionId == null ) { const tagsString = tags . join ( ',' ) ; sessionId = backend . loadSavedModelMetaGraph ( path , tagsString ) ; } const id = nextTFSavedModelId ++ ; const savedModel = new TFSavedModel ( sessionId , id , signatureDefEntry , backend ) ; loadedSavedModelPathMap . set ( id , { path , tags , sessionId } ) ; return savedModel ; } param arrayA param arrayB function stringArraysHaveSameElements ( arrayA , arrayB ) { if ( arrayA . length === arrayB . length && arrayA . sort ( ) . join ( ) === arrayB . sort ( ) . join ( ) ) { return true ; } return false ; } function mapTFDtypeToJSDtype ( tfDtype ) { switch ( tfDtype ) { case 'DT_FLOAT' : return 'float32' ; case 'DT_INT64' : case 'DT_INT32' : case 'DT_UINT8' : return 'int32' ; case 'DT_BOOL' : return 'bool' ; case 'DT_COMPLEX64' : return 'complex64' ; case 'DT_STRING' : return 'string' ; default : throw new Error ( 'Unsupported tensor DataType: ' + tfDtype + ', try to modify the model in python to convert the datatype' ) ; } } export function getNumOfSavedModels ( ) { ensureTensorflowBackend ( ) ; const backend = nodeBackend ( ) ; return backend . getNumOfSavedModels ( ) ; }