license import * as tf from '@tensorflow/tfjs' ; import { backend_util , BackendTimingInfo , DataId , DataType , KernelBackend , ModelTensorInfo , Rank , Scalar , scalar , ScalarLike , Tensor , Tensor1D , Tensor2D , Tensor3D , Tensor4D , TensorInfo , tidy , util } from '@tensorflow/tfjs' ; import { isArray , isNullOrUndefined } from 'util' ; import { encodeInt32ArrayAsInt64 , Int64Scalar } from './int64_tensors' ; import { TensorMetadata , TFEOpAttr , TFJSBinding } from './tfjs_binding' ; const messages = require ( './proto/api_pb' ) ; type TensorData = { shape : number [ ] ; dtype : number ; values : backend_util . BackendValues ; id : number ; refCount : number ; } ; export class NodeJSKernelBackend extends KernelBackend { binding : TFJSBinding ; isGPUPackage : boolean ; isUsingGpuDevice : boolean ; private tensorMap : tf . DataStorage < TensorData > ; constructor ( binding , packageName ) { super ( ) ; this . binding = binding ; this . isGPUPackage = packageName === '@tensorflow/tfjs-node-gpu' ; this . isUsingGpuDevice = this . binding . isUsingGpuDevice ( ) ; this . tensorMap = new tf . DataStorage < TensorData > ( this , tf . engine ( ) ) ; } getDTypeInteger ( dtype ) { switch ( dtype ) { case 'float32' : return this . binding . TF_FLOAT ; case 'int32' : return this . binding . TF_INT32 ; case 'bool' : return this . binding . TF_BOOL ; case 'complex64' : return this . binding . TF_COMPLEX64 ; case 'string' : return this . binding . TF_STRING ; default : throw new Error ( `Unsupported DType: ${ dtype }` ) ; } } private typeAttributeFromTensor ( value ) { return this . getDTypeInteger ( value . dtype ) ; } private createOutputTensor ( metadata ) { const newId = { } ; this . tensorMap . set ( newId , { shape : metadata . shape , dtype : metadata . dtype , id : metadata . id , values : null , refCount : 1 } ) ; let dtype ; switch ( metadata . dtype ) { case this . binding . TF_FLOAT : dtype = 'float32' ; break ; case this . binding . TF_INT32 : dtype = 'int32' ; break ; case this . binding . TF_INT64 : console . warn ( 'INT64 output tensor will be stored as BigInt64Array.' ) ; dtype = 'int32' ; break ; case this . binding . TF_BOOL : dtype = 'bool' ; break ; case this . binding . TF_COMPLEX64 : dtype = 'complex64' ; break ; case this . binding . TF_STRING : dtype = 'string' ; break ; case this . binding . TF_RESOURCE : dtype = 'string' ; break ; case this . binding . TF_UINT8 : dtype = 'int32' ; break ; default : throw new Error ( `Unknown dtype enum ${ metadata . dtype }` ) ; } return tf . engine ( ) . makeTensorFromDataId ( newId , metadata . shape , dtype ) ; } private getInputTensorIds ( tensors ) { const ids = [ ] ; for ( let i = 0 ; i < tensors . length ; i ++ ) { if ( tensors [ i ] instanceof Int64Scalar ) { const value = ( tensors [ i ] as Int64Scalar ) . valueArray ; const id = this . binding . createTensor ( [ ] , this . binding . TF_INT64 , value ) ; ids . push ( id ) ; } else { const info = this . tensorMap . get ( ( tensors [ i ] as TensorInfo ) . dataId ) ; if ( info . values != null ) { info . id = this . binding . createTensor ( info . shape , info . dtype , info . values ) ; info . values = null ; } ids . push ( info . id ) ; } } return ids ; } createReductionOpAttrs ( tensor , keepDims = false ) { return [ { name : 'keep_dims' , type : this . binding . TF_ATTR_BOOL , value : keepDims } , createTensorsTypeOpAttr ( 'T' , tensor . dtype ) , createTensorsTypeOpAttr ( 'Tidx' , 'int32' ) ] ; } floatPrecision ( ) { return 32 ; } epsilon ( ) { return super . epsilon ( ) ; } param name param input executeSingleInput ( name , input ) { const opAttrs = [ createTensorsTypeOpAttr ( 'T' , input . dtype ) ] ; return this . executeSingleOutput ( name , opAttrs , [ input ] ) ; } param name param opAttrs param inputs return executeSingleOutput ( name , opAttrs , inputs ) { const outputMetadata = this . binding . executeOp ( name , opAttrs , this . getInputTensorIds ( inputs ) , 1 ) ; return this . createOutputTensor ( outputMetadata [ 0 ] ) ; } param name param opAttrs param inputs param numOutputs return executeMultipleOutputs ( name , opAttrs , inputs , numOutputs ) { const outputMetadata = this . binding . executeOp ( name , opAttrs , this . getInputTensorIds ( inputs ) , numOutputs ) ; return outputMetadata . map ( m => this . createOutputTensor ( m ) ) ; } numDataIds ( ) { return this . tensorMap . numDataIds ( ) ; } dispose ( ) { } async read ( dataId ) { return this . readSync ( dataId ) ; } readSync ( dataId ) { if ( ! this . tensorMap . has ( dataId ) ) { throw new Error ( `Tensor ${ dataId } was not registered!` ) ; } const info = this . tensorMap . get ( dataId ) ; if ( info . values != null ) { return info . values ; } else { return this . binding . tensorDataSync ( info . id ) ; } } param dataId oaram disposeData ( dataId , force = false ) { if ( this . tensorMap . has ( dataId ) ) { const id = this . tensorMap . get ( dataId ) . id ; this . tensorMap . get ( dataId ) . refCount -- ; if ( ! force && this . tensorMap . get ( dataId ) . refCount > 0 ) { return false ; } if ( id != null && id >= 0 ) { this . binding . deleteTensor ( id ) ; } this . tensorMap . delete ( dataId ) ; } return true ; } /** Return refCount of a `TensorData`. */ refCount ( dataId ) { if ( this . tensorMap . has ( dataId ) ) { const tensorData = this . tensorMap . get ( dataId ) ; return tensorData . refCount ; } return 0 ; } incRef ( dataId ) { this . tensorMap . get ( dataId ) . refCount ++ ; } move ( dataId , values , shape , dtype , refCount ) { this . tensorMap . set ( dataId , { shape , dtype : getTFDType ( dtype ) , values , id : - 1 , refCount } ) ; } write ( values , shape , dtype ) { const dataId = { } ; this . move ( dataId , values , shape , dtype , 1 ) ; return dataId ; } applyActivation < T extends Tensor > ( input , activation , preluActivationWeights ? , leakyreluAlpha ? ) { let result = input ; if ( activation != null ) { if ( activation === 'linear' ) { } else if ( activation === 'relu' ) { result = tf . relu ( result ) ; } else if ( activation === 'prelu' ) { result = tf . prelu ( result , preluActivationWeights ) as T ; } else if ( activation === 'leakyrelu' ) { result = tf . leakyRelu ( result , leakyreluAlpha ) ; } else if ( activation === 'elu' ) { result = tf . elu ( result ) ; } else if ( activation === 'relu6' ) { result = tf . relu6 ( result ) ; } else if ( activation === 'sigmoid' ) { result = tf . sigmoid ( result ) ; } else { throw new Error ( `Activation: ${ activation } has not been implemented for the Node.js backend` ) ; } } return result ; } divide ( a , b ) { const opAttrs = [ createTensorsTypeOpAttr ( 'T' , backend_util . upcastType ( a . dtype , b . dtype ) ) ] ; return this . executeSingleOutput ( 'Div' , opAttrs , [ a , b ] ) ; } divNoNan ( a , b ) { const opAttrs = [ createTensorsTypeOpAttr ( 'T' , backend_util . upcastType ( a . dtype , b . dtype ) ) ] ; return this . executeSingleOutput ( 'DivNoNan' , opAttrs , [ a , b ] ) ; } where ( condition ) { return this . executeSingleOutput ( 'Where' , [ ] , [ condition ] ) as Tensor2D ; } topKValues < T extends Tensor > ( x , k ) { throw new Error ( 'Method not implemented.' ) ; } topKIndices ( x , k ) { throw new Error ( 'Method not implemented.' ) ; } int < T extends Tensor > ( x ) { throw new Error ( 'Method not implemented.' ) ; } decodeJpeg ( contents , channels , ratio , fancyUpscaling , tryRecoverTruncated , acceptableFraction , dctMethod ) { const opAttrs = [ { name : 'channels' , type : this . binding . TF_ATTR_INT , value : channels } , { name : 'ratio' , type : this . binding . TF_ATTR_INT , value : ratio } , { name : 'fancy_upscaling' , type : this . binding . TF_ATTR_BOOL , value : fancyUpscaling } , { name : 'try_recover_truncated' , type : this . binding . TF_ATTR_BOOL , value : tryRecoverTruncated } , { name : 'acceptable_fraction' , type : this . binding . TF_ATTR_FLOAT , value : acceptableFraction } , { name : 'dct_method' , type : this . binding . TF_ATTR_STRING , value : dctMethod } ] ; const inputArgs = [ scalar ( contents , 'string' ) ] ; return this . executeSingleOutput ( 'DecodeJpeg' , opAttrs , inputArgs ) as Tensor < Rank . R3 > ; } decodePng ( contents , channels ) { const opAttrs = [ { name : 'channels' , type : this . binding . TF_ATTR_INT , value : channels } ] ; const inputArgs = [ scalar ( contents , 'string' ) ] ; return this . executeSingleOutput ( 'DecodePng' , opAttrs , inputArgs ) as Tensor < Rank . R3 > ; } decodeBmp ( contents , channels ) { const opAttrs = [ { name : 'channels' , type : this . binding . TF_ATTR_INT , value : channels } ] ; const inputArgs = [ scalar ( contents , 'string' ) ] ; return this . executeSingleOutput ( 'DecodeBmp' , opAttrs , inputArgs ) as Tensor < Rank . R3 > ; } decodeGif ( contents ) { const inputArgs = [ scalar ( contents , 'string' ) ] ; return this . executeSingleOutput ( 'DecodeGif' , [ ] , inputArgs ) as Tensor < Rank . R4 > ; } executeEncodeImageOp ( name , opAttrs , imageData , imageShape ) { const inputTensorId = this . binding . createTensor ( imageShape , this . binding . TF_UINT8 , imageData ) ; const outputMetadata = this . binding . executeOp ( name , opAttrs , [ inputTensorId ] , 1 ) ; const outputTensorInfo = outputMetadata [ 0 ] ; outputTensorInfo . dtype = this . binding . TF_UINT8 ; return this . createOutputTensor ( outputTensorInfo ) ; } encodeJpeg ( imageData , imageShape , format , quality , progressive , optimizeSize , chromaDownsampling , densityUnit , xDensity , yDensity , xmpMetadata ) { const opAttrs = [ { name : 'format' , type : this . binding . TF_ATTR_STRING , value : format } , { name : 'quality' , type : this . binding . TF_ATTR_INT , value : quality } , { name : 'progressive' , type : this . binding . TF_ATTR_BOOL , value : progressive } , { name : 'optimize_size' , type : this . binding . TF_ATTR_BOOL , value : optimizeSize } , { name : 'chroma_downsampling' , type : this . binding . TF_ATTR_BOOL , value : chromaDownsampling } , { name : 'density_unit' , type : this . binding . TF_ATTR_STRING , value : densityUnit } , { name : 'x_density' , type : this . binding . TF_ATTR_INT , value : xDensity } , { name : 'y_density' , type : this . binding . TF_ATTR_INT , value : yDensity } , { name : 'xmp_metadata' , type : this . binding . TF_ATTR_STRING , value : xmpMetadata } ] ; return this . executeEncodeImageOp ( 'EncodeJpeg' , opAttrs , imageData , imageShape ) ; } encodePng ( imageData , imageShape , compression ) { const opAttrs = [ { name : 'compression' , type : this . binding . TF_ATTR_INT , value : compression } ] ; return this . executeEncodeImageOp ( 'EncodePng' , opAttrs , imageData , imageShape ) ; } deleteSavedModel ( id ) { this . binding . deleteSavedModel ( id ) ; } loadSavedModelMetaGraph ( path , tags ) { return this . binding . loadSavedModel ( path , tags ) ; } private getMappedInputTensorIds ( inputs , inputTensorInfos ) { const tensorIds = this . getInputTensorIds ( inputs ) ; for ( let i = 0 ; i < inputs . length ; i ++ ) { if ( inputTensorInfos [ i ] != null ) { if ( 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 ; } else if ( inputTensorInfos [ i ] . tfDtype === 'DT_INT64' ) { const data = encodeInt32ArrayAsInt64 ( inputs [ i ] . dataSync ( ) as Int32Array ) ; const inputTensorId = this . binding . createTensor ( inputs [ i ] . shape , this . binding . TF_INT64 , data ) ; tensorIds [ i ] = inputTensorId ; } } } return tensorIds ; } runSavedModel ( id , inputs , inputTensorInfos , outputOpNames ) { const outputMetadata = this . binding . runSavedModel ( id , this . getMappedInputTensorIds ( inputs , inputTensorInfos ) , inputTensorInfos . map ( info => info . name ) . join ( ',' ) , outputOpNames . join ( ',' ) ) ; return outputMetadata . map ( m => this . createOutputTensor ( m ) ) ; } summaryWriter ( logdir ) { const opAttrs = [ { name : 'shared_name' , type : this . binding . TF_ATTR_STRING , value : `logdir:${ logdir }` } , { name : 'container' , type : this . binding . TF_ATTR_STRING , value : '' } ] ; const writerResource = this . executeSingleOutput ( 'SummaryWriter' , opAttrs , [ ] ) ; return writerResource as Tensor1D ; } createSummaryFileWriter ( resourceHandle , logdir , maxQueue ? , flushMillis ? , filenameSuffix ? ) { const inputArgs = [ resourceHandle , scalar ( logdir ) , scalar ( maxQueue == null ? 10 : maxQueue , 'int32' ) , scalar ( flushMillis == null ? 2 * 60 * 1000 : flushMillis , 'int32' ) , scalar ( filenameSuffix == null ? '.v2' : filenameSuffix ) ] ; this . executeMultipleOutputs ( 'CreateSummaryFileWriter' , [ ] , inputArgs , 0 ) ; } writeScalarSummary ( resourceHandle , step , name , value ) { tidy ( ( ) => { util . assert ( Number . isInteger ( step ) , ( ) => `step is expected to be an integer, but is instead ${ step }` ) ; const inputArgs = [ resourceHandle , new Int64Scalar ( step ) , scalar ( name , 'string' ) ] ; let typeAttr ; if ( typeof value === 'number' ) { inputArgs . push ( scalar ( value ) ) ; typeAttr = this . binding . TF_FLOAT ; } else { util . assert ( value . rank === 0 , ( ) => `A non-scalar tensor (rank ${ value . rank }) is passed to ` + `writeScalarSummary()` ) ; inputArgs . push ( value ) ; typeAttr = this . typeAttributeFromTensor ( value ) ; } const opAttrs = [ { name : 'T' , type : this . binding . TF_ATTR_TYPE , value : typeAttr } ] ; this . binding . executeOp ( 'WriteScalarSummary' , opAttrs , this . getInputTensorIds ( inputArgs ) , 0 ) ; } ) ; } writeHistogramSummary ( resourceHandle , step , name , data , bucketCount , description ) { tidy ( ( ) => { util . assert ( Number . isInteger ( step ) , ( ) => `step is expected to be an integer, but is instead ${ step }` ) ; const content = new messages . HistogramPluginData ( ) . setVersion ( 0 ) ; const pluginData = new messages . SummaryMetadata . PluginData ( ) . setPluginName ( 'histograms' ) . setContent ( content . serializeBinary ( ) ) ; const summary = new messages . SummaryMetadata ( ) . setPluginData ( pluginData ) . setDisplayName ( null ) . setSummaryDescription ( description ) ; const summaryTensor = scalar ( summary . serializeBinary ( ) , 'string' ) ; const nameTensor = scalar ( name , 'string' ) ; const stepScalar = new Int64Scalar ( step ) ; const buckets = this . buckets ( data , bucketCount ) ; util . assert ( buckets . rank === 2 && buckets . shape [ 1 ] === 3 , ( ) => `Expected buckets to have shape [k, 3], but they had shape ${ buckets . shape }` ) ; util . assert ( buckets . dtype === 'float32' , ( ) => `Expected buckets to have dtype float32, but they had dtype ${ buckets . dtype }` ) ; const inputArgs = [ resourceHandle , stepScalar , buckets , nameTensor , summaryTensor ] ; const typeAttr = this . typeAttributeFromTensor ( buckets ) ; const opAttrs = [ { name : 'T' , type : this . binding . TF_ATTR_TYPE , value : typeAttr } ] ; this . binding . executeOp ( 'WriteSummary' , opAttrs , this . getInputTensorIds ( inputArgs ) , 0 ) ; } ) ; } flushSummaryWriter ( resourceHandle ) { const inputArgs = [ resourceHandle ] ; this . executeMultipleOutputs ( 'FlushSummaryWriter' , [ ] , inputArgs , 0 ) ; } param data param bucketCount returns private buckets ( data , bucketCount ? ) { if ( data . size === 0 ) { return tf . tensor ( [ ] , [ 0 , 3 ] , 'float32' ) ; } bucketCount = bucketCount !== undefined ? bucketCount : 30 ; util . assert ( Number . isInteger ( bucketCount ) && bucketCount > 0 , ( ) => `Expected bucket count to be a strictly positive integer, but it was ` + `${ bucketCount }` ) ; data = data . flatten ( ) ; data = data . cast ( 'float32' ) ; const min = data . min ( ) ; const max = data . max ( ) ; const range = max . sub ( min ) ; const isSingular = range . equal ( 0 ) . arraySync ( ) !== 0 ; if ( isSingular ) { const center = min ; const bucketStart = center . sub ( 0.5 ) ; const bucketEnd = center . add ( 0.5 ) ; const bucketCounts = tf . scalar ( data . size , 'float32' ) ; return tf . concat ( [ bucketStart , bucketEnd , bucketCounts ] ) . reshape ( [ 1 , 3 ] ) ; } const bucketWidth = range . div ( bucketCount ) ; const offsets = data . sub ( min ) ; const bucketIndices = offsets . floorDiv ( bucketWidth ) . cast ( 'int32' ) ; const clampedIndices = tf . minimum ( bucketIndices , bucketCount - 1 ) . cast ( 'int32' ) ; const oneHots = tf . oneHot ( clampedIndices , bucketCount ) ; const bucketCounts = oneHots . sum ( 0 ) . cast ( 'int32' ) ; let edges = tf . linspace ( min . arraySync ( ) , max . arraySync ( ) , bucketCount + 1 ) ; edges = tf . concat ( [ edges . slice ( 0 , bucketCount ) , max . reshape ( [ 1 ] ) ] , 0 ) as tf . Tensor1D ; const leftEdges = edges . slice ( 0 , bucketCount ) ; const rightEdges = edges . slice ( 1 , bucketCount ) ; return tf . stack ( [ leftEdges , rightEdges , bucketCounts . cast ( 'float32' ) ] ) . transpose ( ) ; } memory ( ) { return { unreliable : true } ; } async time ( f ) { const start = process . hrtime ( ) ; f ( ) ; const elapsed = process . hrtime ( start ) ; return { kernelMs : elapsed [ 0 ] * 1000 + elapsed [ 1 ] / 1000000 } ; } getNumOfSavedModels ( ) { return this . binding . getNumOfSavedModels ( ) ; } } /** Returns an instance of the Node.js backend. */ export function nodeBackend ( ) { return tf . findBackend ( 'tensorflow' ) as NodeJSKernelBackend ; } /** Returns the TF dtype for a given DataType. */ export function getTFDType ( dataType ) { const binding = nodeBackend ( ) . binding ; switch ( dataType ) { case 'float32' : return binding . TF_FLOAT ; case 'int32' : return binding . TF_INT32 ; case 'bool' : return binding . TF_BOOL ; case 'complex64' : return binding . TF_COMPLEX64 ; case 'string' : return binding . TF_STRING ; case 'int64' as any : return binding . TF_INT64 ; default : const errorMessage = `Unknown dtype: ${ dataType }` ; throw new Error ( errorMessage ) ; } } /**
 * Creates a TFEOpAttr for a 'type' OpDef attribute from a Tensor or list of
 * Tensors.
 */ export function createTensorsTypeOpAttr ( attrName , tensorsOrDtype ) { if ( isNullOrUndefined ( tensorsOrDtype ) ) { throw new Error ( 'Invalid input tensors value.' ) ; } return { name : attrName , type : nodeBackend ( ) . binding . TF_ATTR_TYPE , value : ( tensorsOrDtype instanceof tf . Tensor || Array . isArray ( tensorsOrDtype ) ) ? getTFDTypeForInputs ( tensorsOrDtype ) : getTFDType ( tensorsOrDtype ) } ; } export function createOpAttr ( attrName , tensorsOrDtype , value ) { if ( isNullOrUndefined ( tensorsOrDtype ) ) { throw new Error ( 'Invalid input tensors value.' ) ; } return { name : attrName , type : nodeBackend ( ) . binding . TF_BOOL , value } ; } /** Returns the dtype number for a single or list of input Tensors. */ function getTFDTypeForInputs ( tensors ) { if ( isNullOrUndefined ( tensors ) ) { throw new Error ( 'Invalid input tensors value.' ) ; } if ( isArray ( tensors ) ) { for ( let i = 0 ; i < tensors . length ; i ++ ) { return getTFDType ( tensors [ i ] . dtype ) ; } return - 1 ; } else { return getTFDType ( tensors . dtype ) ; } } export function ensureTensorflowBackend ( ) { tf . util . assert ( tf . getBackend ( ) === 'tensorflow' , ( ) => `Expect the current backend to be "tensorflow", but got "${ tf . getBackend ( ) }"` ) ; }