license /**
 *  This file is used to load a saved model and perform inference.
 *  Run this script in console:
 *   ts-node index.ts --model_path=MODEL_PATH --inputs_dir=INPUTS_DIR
 *   --outputs_dir=OUTPUTS_DIR
 *
 *  For help, run:
 *   ts-node inference.ts -h
 */ import '@tensorflow/tfjs-backend-wasm' ; import '@tensorflow/tfjs-backend-cpu' ; import * as tfconv from '@tensorflow/tfjs-converter' ; import * as tfc from '@tensorflow/tfjs-core' ; import * as fs from 'fs' ; import * as path from 'path' ; import * as yargs from 'yargs' ; import { FileHandler } from './file_handler' ; interface Options { model_path : string ; inputs_dir : string ; outputs_dir : string ; inputs_data_file : string ; inputs_shape_file : string ; inputs_dtype_file : string ; tf_input_name_file : string ; tf_output_name_file : string ; backend : string ; } /**
 *   The entry point of tfjs inference api.
 */ async function main ( ) { const argParser = yargs . options ( { model_path : { description : 'Path to the tfjs model json file.' , type : 'string' , demandOption : true } , inputs_dir : { description : 'Directory to read the input tensor info and output ' + 'info files.' , type : 'string' , demandOption : true } , outputs_dir : { description : 'Directory to write the output files. Output files include: ' + 'data.json, shape.json and dtype.json. The order of the output ' + 'tensors follow the same order as the tf_output_name_file. If the ' + 'file is not provided, the default outputs of the model would be ' + 'used.' , type : 'string' , demandOption : true } , inputs_data_file : { description : 'Filename of the input data file. Must provide the file.' , type : 'string' , default : 'data.json' } , inputs_shape_file : { description : 'Filename of the input shape file. Must provide the file.' , type : 'string' , default : 'shape.json' } , inputs_dtype_file : { description : 'Filename of the input dtype file. Must provide the file.' , type : 'string' , default : 'dtype.json' } , tf_input_name_file : { description : 'Filename of the input name of the tf model. The input ' + 'names should match the names defined in the signatureDef of the ' + 'model. Must provide the file.' , type : 'string' , default : 'tf_input_name.json' } , tf_output_name_file : { description : 'Filename of the output name of the tf model. The output ' + 'names should match the names defined in the signatureDef of the ' + 'model. If the file is not provided, the default outputs of the ' + 'model would be used.' , type : 'string' } , backend : { description : 'Choose which tfjs backend to use. Supported backends: ' + 'cpu|wasm' , type : 'string' , default : 'cpu' } } ) ; const options = argParser . argv as { } as Options ; if ( options . backend === 'wasm' ) { await tfc . setBackend ( 'wasm' ) ; } else if ( options . backend === 'cpu' ) { await tfc . setBackend ( 'cpu' ) ; } else { throw new Error ( 'Only cpu and wasm backend is supported, but got ' + options . backend ) ; } const model = await tfconv . loadGraphModel ( new FileHandler ( options . model_path ) ) ; const inputsDataString = fs . readFileSync ( path . join ( options . inputs_dir , options . inputs_data_file ) , 'utf8' ) ; const inputsData = JSON . parse ( inputsDataString ) ; const inputsShapeString = fs . readFileSync ( path . join ( options . inputs_dir , options . inputs_shape_file ) , 'utf8' ) ; const inputsShape = JSON . parse ( inputsShapeString ) ; const inputsDtypeString = fs . readFileSync ( path . join ( options . inputs_dir , options . inputs_dtype_file ) , 'utf8' ) ; const inputsDtype = JSON . parse ( inputsDtypeString ) ; const tfInputNameString = fs . readFileSync ( path . join ( options . inputs_dir , options . tf_input_name_file ) , 'utf8' ) ; const inputName = JSON . parse ( tfInputNameString ) ; let outputName ; if ( options . tf_output_name_file ) { const tfOutputNameString = fs . readFileSync ( path . join ( options . inputs_dir , options . tf_output_name_file ) , 'utf8' ) ; outputName = JSON . parse ( tfOutputNameString ) ; } const namedInputs = createInputTensors ( inputsData , inputsShape , inputsDtype , inputName ) ; const result = await model . executeAsync ( namedInputs , outputName ) ; const ys = ( model . outputs . length === 1 ? [ result ] : result ) as tfc . Tensor [ ] ; const ysData = [ ] ; const ysShape = [ ] ; const ysDtype = [ ] ; for ( let i = 0 ; i < ys . length ; i ++ ) { const y = ys [ i ] ; ysData . push ( await y . data ( ) ) ; ysShape . push ( y . shape ) ; ysDtype . push ( y . dtype ) ; } fs . writeFileSync ( path . join ( options . outputs_dir , 'data.json' ) , JSON . stringify ( ysData ) ) ; fs . writeFileSync ( path . join ( options . outputs_dir , 'shape.json' ) , JSON . stringify ( ysShape ) ) ; fs . writeFileSync ( path . join ( options . outputs_dir , 'dtype.json' ) , JSON . stringify ( ysDtype ) ) ; Object . keys ( namedInputs ) . forEach ( key => namedInputs [ key ] . dispose ( ) ) ; ys . forEach ( tensor => tensor . dispose ( ) ) ; } private param inputsData param inputsShape param inputsDtype param inputName return function createInputTensors ( inputsData , inputsShape , inputsDtype , inputName ) { const xs = [ ] ; for ( let i = 0 ; i < inputsData . length ; i ++ ) { const input = tfc . tensor ( inputsData [ i ] , inputsShape [ i ] , inputsDtype [ i ] ) ; xs . push ( input ) ; } return inputName . reduce ( ( map , name , index ) => { map [ name ] = xs [ index ] ; return map ; } , { } ) ; } main ( ) ;