license import { DataType , env , NamedTensorMap , Tensor , tidy , util } from '@tensorflow/tfjs-core' ; import { ISignatureDef } from '../data/compiled_api' ; import { NamedTensorsMap , TensorArrayMap , TensorInfo , TensorListMap } from '../data/types' ; import { getNodeNameAndIndex , getParamValue , getTensor , getTensorsForCurrentContenxt , parseNodeName } from '../operations/executors/utils' ; import { executeOp } from '../operations/operation_executor' ; import { Graph , Node } from '../operations/types' ; import { ExecutionContext , ExecutionContextInfo } from './execution_context' ; import { getExecutionSubgraph , getNodesInTopologicalOrder , isControlFlow } from './model_analysis' ; import { ResourceManager } from './resource_manager' ; import { FunctionExecutor } from './types' ; interface NodeWithContexts { contexts : ExecutionContextInfo [ ] ; node : Node ; } export class GraphExecutor implements FunctionExecutor { private compiledMap : Map < string , Node [ ] > = new Map ( ) ; private _weightMap : NamedTensorsMap = { } ; private _weightIds : number [ ] ; private _signature : ISignatureDef ; private _inputs : Node [ ] ; private _outputs : Node [ ] ; private _initNodes : Node [ ] ; private SEPERATOR = ',' ; private _functions : { [ key ] : Graph ; } = { } ; private _functionExecutorMap : { [ key ] : FunctionExecutor ; } = { } ; private _resourceManager : ResourceManager ; private intermediateTensors : NamedTensorsMap = { } ; private keepIds : Set < number > ; private tensorsMap : NamedTensorsMap ; private keepTensorForDebug = false ; get weightIds ( ) : number [ ] { return this . parent ? this . parent . weightIds : this . _weightIds ; } get functionExecutorMap ( ) : { [ key ] : FunctionExecutor ; } { return this . parent ? this . parent . functionExecutorMap : this . _functionExecutorMap ; } get weightMap ( ) : NamedTensorsMap { return this . parent ? this . parent . weightMap : this . _weightMap ; } set weightMap ( weightMap ) { const weightIds = Object . keys ( weightMap ) . map ( key => weightMap [ key ] . map ( tensor => tensor . id ) ) ; this . _weightIds = [ ] . concat ( ... weightIds ) ; this . _weightMap = weightMap ; } param resourceManager set resourceManager ( resourceManager ) { this . _resourceManager = resourceManager ; } get inputs ( ) : TensorInfo [ ] { return this . _inputs . map ( node => { return { name : node . name , shape : node . attrParams [ 'shape' ] ? node . attrParams [ 'shape' ] . value as number [ ] : undefined , dtype : node . attrParams [ 'dtype' ] ? node . attrParams [ 'dtype' ] . value as DataType : undefined } ; } ) ; } get outputs ( ) : TensorInfo [ ] { return this . _outputs . map ( node => { return { name : node . name , shape : node . attrParams [ 'shape' ] ? node . attrParams [ 'shape' ] . value as number [ ] : undefined , dtype : node . attrParams [ 'dtype' ] ? node . attrParams [ 'dtype' ] . value as DataType : undefined } ; } ) ; } get inputNodes ( ) : string [ ] { return this . _inputs . map ( node => node . signatureKey || node . name ) ; } get outputNodes ( ) : string [ ] { return this . _outputs . map ( ( node ) => { const name = node . signatureKey || node . name ; return node . defaultOutput ? ( `${ name }:${ node . defaultOutput }` ) : name ; } ) ; } get functions ( ) : { [ key ] : ISignatureDef ; } { return Object . keys ( this . _functions ) . reduce ( ( map , key ) => { map [ key ] = this . _functions [ key ] . signature ; return map ; } , { } as { [ key ] : ISignatureDef ; } ) ; } param graph param parent constructor ( private graph , private parent ? ) { this . _outputs = graph . outputs ; this . _inputs = graph . inputs ; this . _initNodes = graph . initNodes ; this . _signature = graph . signature ; this . _functions = graph . functions ; if ( graph . functions != null ) { Object . keys ( graph . functions ) . forEach ( name => { this . _functionExecutorMap [ name ] = new GraphExecutor ( graph . functions [ name ] , this ) ; } ) ; } } private getCompilationKey ( inputs , outputs ) { const sortedInputs = inputs . map ( node => node . name ) . sort ( ) ; const sortedOutputs = outputs . map ( node => node . name ) . sort ( ) ; return sortedInputs . join ( this . SEPERATOR ) + '--' + sortedOutputs . join ( this . SEPERATOR ) ; } /**
     * Compiles the inference graph and returns the minimal set of nodes that are
     * required for execution, in the correct execution order.
     */ private compile ( inputs , outputs ) { const executionInfo = getExecutionSubgraph ( inputs , outputs , this . weightMap , this . _initNodes ) ; const { missingInputs , dynamicNode , syncInputs } = executionInfo ; if ( dynamicNode != null ) { throw new Error ( `This execution contains the node '${ dynamicNode . name }', which has ` + `the dynamic op '${ dynamicNode . op }'. Please use ` + `model.executeAsync() instead. Alternatively, to avoid the ` + `dynamic ops, specify the inputs [${ syncInputs }]` ) ; } if ( missingInputs . length > 0 ) { const outNames = outputs . map ( n => n . name ) ; const inNames = Object . keys ( inputs ) ; throw new Error ( `Cannot compute the outputs [${ outNames }] from the provided inputs ` + `[${ inNames }]. Missing the following inputs: [${ missingInputs }]` ) ; } return getNodesInTopologicalOrder ( this . graph , this . weightMap , executionInfo ) ; } param inputs param outputs execute ( inputs , outputs ? ) { inputs = this . mapInputs ( inputs ) ; const names = Object . keys ( inputs ) . sort ( ) ; this . checkInputs ( inputs ) ; this . checkInputShapeAndType ( inputs ) ; outputs = this . mapOutputs ( outputs ) ; this . checkOutputs ( outputs ) ; const inputNodes = names . map ( name => this . graph . nodes [ parseNodeName ( name ) [ 0 ] ] ) ; const outputNodeNames = outputs . map ( name => parseNodeName ( name ) [ 0 ] ) ; let outputNodes = outputNodeNames . map ( name => this . graph . nodes [ name ] ) ; this . resetIntermediateTensors ( ) ; if ( outputNodes . length === 0 ) { outputNodes = this . _outputs ; } const compilationKey = this . getCompilationKey ( inputNodes , outputNodes ) ; let orderedNodes = this . compiledMap . get ( compilationKey ) ; if ( orderedNodes == null ) { orderedNodes = this . compile ( inputs , outputNodes ) ; this . compiledMap . set ( compilationKey , orderedNodes ) ; } const tensorArrayMap = { } ; const tensorListMap = { } ; return tidy ( ( ) => { const context = new ExecutionContext ( this . weightMap , tensorArrayMap , tensorListMap , this . functionExecutorMap ) ; const tensorsMap = { ... this . weightMap } ; Object . keys ( inputs ) . forEach ( name => { const [ nodeName , index ] = parseNodeName ( name ) ; const tensors = [ ] ; tensors [ index ] = inputs [ name ] ; tensorsMap [ nodeName ] = tensors ; } ) ; const tensorsToKeep = this . getFrozenTensorIds ( tensorsMap ) ; const intermediateTensorConsumerCount = { } ; for ( let i = 0 ; i < orderedNodes . length ; i ++ ) { const node = orderedNodes [ i ] ; if ( ! tensorsMap [ node . name ] ) { const tensors = executeOp ( node , tensorsMap , context , this . _resourceManager ) as Tensor [ ] ; if ( util . isPromise ( tensors ) ) { throw new Error ( `The execution of the op '${ node . op }' returned a promise. ` + `Please use model.executeAsync() instead.` ) ; } tensorsMap [ node . name ] = tensors ; this . checkTensorForDisposal ( node . name , node , tensorsMap , context , tensorsToKeep , outputNodeNames , intermediateTensorConsumerCount ) ; } } if ( this . parent == null ) { context . dispose ( tensorsToKeep ) ; } return outputs . map ( name => getTensor ( name , tensorsMap , context ) ) ; } ) ; } private getFrozenTensorIds ( tensorMap ) { const ids = [ ] . concat . apply ( [ ] , Object . keys ( tensorMap ) . map ( key => tensorMap [ key ] ) . map ( tensors => tensors . map ( tensor => tensor . id ) ) ) ; return new Set ( ids ) ; } private checkTensorForDisposal ( nodeName , node , tensorMap , context , tensorsToKeep , outputNames , intermediateTensorConsumerCount ) { if ( node . category === 'control' || outputNames . indexOf ( nodeName ) !== - 1 ) { return ; } tensorMap [ nodeName ] . forEach ( tensor => { if ( tensor != null ) { intermediateTensorConsumerCount [ tensor . id ] = ( intermediateTensorConsumerCount [ tensor . id ] || 0 ) + node . children . length ; } } ) ; node . inputs . forEach ( input => { if ( input . category !== 'control' ) { const tensors = getTensorsForCurrentContenxt ( input . name , tensorMap , context ) ; if ( tensors != null ) { tensors . forEach ( tensor => { if ( tensor && ! tensor . kept && ! tensorsToKeep . has ( tensor . id ) ) { const count = intermediateTensorConsumerCount [ tensor . id ] ; if ( count === 1 ) { if ( ! this . keepTensorForDebug ) { tensor . dispose ( ) ; } else { const [ nodeName , index ] = getNodeNameAndIndex ( node . name , context ) ; if ( this . intermediateTensors [ nodeName ] ) { this . intermediateTensors [ nodeName ] [ index ] = tensor ; } else { this . intermediateTensors [ nodeName ] = [ ] ; this . intermediateTensors [ nodeName ] [ index ] = tensor ; } } delete intermediateTensorConsumerCount [ tensor . id ] ; } else if ( count != null ) { intermediateTensorConsumerCount [ tensor . id ] -- ; } } } ) ; } } } ) ; } param inputs param outputs async executeAsync ( inputs , outputs ? ) { return this . _executeAsync ( inputs , outputs ) ; } disposeIntermediateTensors ( ) { if ( ! this . intermediateTensors ) { return ; } Object . keys ( this . intermediateTensors ) . forEach ( key => this . intermediateTensors [ key ] . forEach ( tensor => tensor . dispose ( ) ) ) ; this . disposeTensorsMap ( ) ; } private disposeTensorsMap ( ) { if ( ! this . tensorsMap ) { return ; } Object . keys ( this . tensorsMap ) . forEach ( key => { const tensorArray = this . tensorsMap [ key ] ; tensorArray . forEach ( tensor => { if ( tensor && ! tensor . kept && ! tensor . isDisposed && ! this . keepIds . has ( tensor . id ) ) { tensor . dispose ( ) ; } } ) ; } ) ; } getIntermediateTensors ( ) { return this . tensorsMap ; } private resetIntermediateTensors ( ) { for ( const key in this . intermediateTensors ) { this . intermediateTensors [ key ] . forEach ( tensor => tensor . dispose ( ) ) ; delete this . intermediateTensors [ key ] ; } } param inputs param outputs param isFunctionExecution param tensorArrayMap param tensorArrayMap private async _executeAsync ( inputs , outputs ? , isFunctionExecution = false , tensorArrayMap = { } , tensorListMap = { } ) { if ( ! isFunctionExecution ) { inputs = this . mapInputs ( inputs ) ; this . checkInputs ( inputs ) ; this . checkInputShapeAndType ( inputs ) ; outputs = this . mapOutputs ( outputs ) ; this . checkOutputs ( outputs ) ; } try { this . keepTensorForDebug = env ( ) . getBool ( 'KEEP_INTERMEDIATE_TENSORS' ) ; } catch ( e ) { console . warn ( e . message ) ; } this . resetIntermediateTensors ( ) ; const context = new ExecutionContext ( this . weightMap , tensorArrayMap , tensorListMap , this . functionExecutorMap ) ; this . tensorsMap = await this . executeWithControlFlow ( inputs , context , outputs , isFunctionExecution ) ; const results = outputs . map ( name => getTensor ( name , this . tensorsMap , context ) ) ; const outputIds = results . map ( t => t . id ) ; const inputIds = Object . keys ( inputs ) . map ( name => inputs [ name ] . id ) ; this . keepIds = new Set < number > ( [ ... outputIds , ... inputIds , ... this . weightIds ] ) ; if ( ! this . keepTensorForDebug ) { this . disposeTensorsMap ( ) ; } if ( this . parent == null ) { context . dispose ( this . keepIds ) ; } return results ; } async executeFunctionAsync ( inputs , tensorArrayMap , tensorListMap ) { const mappedInputs = inputs . reduce ( ( map , tensor , index ) => { map [ this . inputs [ index ] . name ] = tensor ; return map ; } , { } as NamedTensorMap ) ; return this . _executeAsync ( mappedInputs , this . outputNodes , true , tensorArrayMap , tensorListMap ) ; } param inputs param context param outputNames param isFunctionExecution private async executeWithControlFlow ( inputs , context , outputNames ? , isFunctionExecution ? ) { const names = Object . keys ( inputs ) ; const inputNodes = names . map ( name => this . graph . nodes [ parseNodeName ( name ) [ 0 ] ] ) ; const outputNodeNames = outputNames . map ( name => parseNodeName ( name ) [ 0 ] ) ; let outputNodes = outputNodeNames . map ( name => this . graph . nodes [ name ] ) ; if ( outputNodes . length === 0 ) { outputNodes = this . _outputs ; } const { usedNodes , missingInputs , dynamicNode , syncInputs } = getExecutionSubgraph ( inputs , outputNodes , this . weightMap , this . _initNodes ) ; const stack = [ ... inputNodes , ... this . graph . weights , ... ( this . _initNodes || [ ] ) ] . map ( node => { return { node , contexts : context . currentContext } ; } ) ; const tensorsMap = { ... this . weightMap } ; Object . keys ( inputs ) . forEach ( name => { const [ nodeName , index ] = parseNodeName ( name ) ; const tensors = [ ] ; tensors [ index ] = inputs [ name ] ; tensorsMap [ nodeName ] = tensors ; } ) ; const intermediateTensorConsumerCount = { } ; const tensorsToKeep = this . getFrozenTensorIds ( tensorsMap ) ; const added = { } ; while ( stack . length > 0 ) { const promises = this . processStack ( inputNodes , stack , context , tensorsMap , added , tensorsToKeep , outputNodeNames , intermediateTensorConsumerCount , usedNodes ) ; await Promise . all ( promises ) ; } if ( dynamicNode == null && ! isFunctionExecution ) { console . warn ( `This model execution did not contain any nodes with control flow ` + `or dynamic output shapes. You can use model.execute() instead.` ) ; } const missingOutputs = outputNodes . filter ( node => ! isControlFlow ( node ) && ! getTensor ( node . name , tensorsMap , context ) ) . map ( node => node . name ) ; if ( missingOutputs . length > 0 ) { let alternativeMsg = '' ; if ( dynamicNode != null ) { alternativeMsg = `Alternatively, to avoid the dynamic ops, use model.execute() ` + `and specify the inputs [${ syncInputs }]` ; } throw new Error ( `Cannot compute the outputs [${ missingOutputs }] from the provided ` + `inputs [${ names }]. Consider providing the following inputs: ` + `[${ missingInputs }]. ${ alternativeMsg }` ) ; } return tensorsMap ; } private processStack ( inputNodes , stack , context , tensorMap , added , tensorsToKeep , outputNames , intermediateTensorConsumerCount , usedNodes ) { const promises = [ ] ; while ( stack . length > 0 ) { const item = stack . pop ( ) ; context . currentContext = item . contexts ; let nodeName = '' ; if ( item . node . op === 'Enter' && getParamValue ( 'isConstant' , item . node , tensorMap , context ) ) { [ nodeName ] = getNodeNameAndIndex ( item . node . name , context ) ; } if ( tensorMap [ item . node . name ] == null ) { const tensors = executeOp ( item . node , tensorMap , context , this . _resourceManager ) ; if ( ! nodeName ) { [ nodeName ] = getNodeNameAndIndex ( item . node . name , context ) ; } const currentContext = context . currentContext ; if ( util . isPromise ( tensors ) ) { promises . push ( tensors . then ( t => { tensorMap [ nodeName ] = t ; context . currentContext = currentContext ; this . checkTensorForDisposal ( nodeName , item . node , tensorMap , context , tensorsToKeep , outputNames , intermediateTensorConsumerCount ) ; this . processChildNodes ( item . node , stack , context , tensorMap , added , usedNodes ) ; return t ; } ) ) ; } else { tensorMap [ nodeName ] = tensors ; this . checkTensorForDisposal ( nodeName , item . node , tensorMap , context , tensorsToKeep , outputNames , intermediateTensorConsumerCount ) ; this . processChildNodes ( item . node , stack , context , tensorMap , added , usedNodes ) ; } } else { this . processChildNodes ( item . node , stack , context , tensorMap , added , usedNodes ) ; } } return promises ; } private processChildNodes ( node , stack , context , tensorMap , added , usedNodes ) { node . children . forEach ( ( childNode ) => { const [ nodeName , ] = getNodeNameAndIndex ( childNode . name , context ) ; if ( added [ nodeName ] || ! usedNodes . has ( childNode . name ) ) { return ; } if ( childNode . op === 'Merge' ) { if ( childNode . inputNames . some ( name => { return ! ! getTensor ( name , tensorMap , context ) ; } ) ) { added [ nodeName ] = true ; stack . push ( { contexts : context . currentContext , node : childNode } ) ; } } else if ( childNode . inputNames . every ( name => { return ! ! getTensor ( name , tensorMap , context ) ; } ) ) { added [ nodeName ] = true ; stack . push ( { contexts : context . currentContext , node : childNode } ) ; } } ) ; } /**
     * Releases the memory used by the weight tensors.
     */ dispose ( ) { Object . keys ( this . weightMap ) . forEach ( key => this . weightMap [ key ] . forEach ( tensor => tensor . dispose ( ) ) ) ; } private checkInputShapeAndType ( inputs ) { Object . keys ( inputs ) . forEach ( name => { const input = inputs [ name ] ; const [ nodeName , ] = parseNodeName ( name ) ; const node = this . graph . nodes [ nodeName ] ; if ( node . attrParams [ 'shape' ] && node . attrParams [ 'shape' ] . value ) { const shape = node . attrParams [ 'shape' ] . value as number [ ] ; const match = shape . length === input . shape . length && input . shape . every ( ( dim , index ) => shape [ index ] === - 1 || shape [ index ] === dim ) ; util . assert ( match , ( ) => `The shape of dict['${ node . name }'] provided in ` + `model.execute(dict) must be [${ shape }], but was ` + `[${ input . shape }]` ) ; } if ( node . attrParams [ 'dtype' ] && node . attrParams [ 'dtype' ] . value ) { util . assert ( input . dtype === node . attrParams [ 'dtype' ] . value as string , ( ) => `The dtype of dict['${ node . name }'] provided in ` + `model.execute(dict) must be ` + `${ node . attrParams [ 'dtype' ] . value }, but was ${ input . dtype }` ) ; } } ) ; } private mapInputs ( inputs ) { const result = { } ; for ( const inputName in inputs ) { if ( this . _signature != null && this . _signature . inputs != null && this . _signature . inputs [ inputName ] != null ) { const tensor = this . _signature . inputs [ inputName ] ; result [ tensor . name ] = inputs [ inputName ] ; } else { result [ inputName ] = inputs [ inputName ] ; } } return result ; } private checkInputs ( inputs ) { const notInGraph = Object . keys ( inputs ) . filter ( name => { const [ nodeName ] = parseNodeName ( name ) ; return this . graph . nodes [ nodeName ] == null ; } ) ; if ( notInGraph . length > 0 ) { throw new Error ( `The dict provided in model.execute(dict) has ` + `keys: [${ notInGraph }] that are not part of graph` ) ; } } private mapOutputs ( outputs ) { return outputs . map ( name => { if ( this . _signature != null && this . _signature . outputs != null && this . _signature . outputs [ name ] != null ) { const tensor = this . _signature . outputs [ name ] ; return tensor . name ; } return name ; } , { } ) ; } private checkOutputs ( outputs ) { outputs . forEach ( name => { const [ normalizedName ] = parseNodeName ( name ) ; if ( ! this . graph . nodes [ normalizedName ] ) { throw new Error ( `The output '${ name }' is not found in the graph` ) ; } } ) ; } }