use tidy() in prediction preprocess to GC intermediate tensor (#122)

tensorspace-team · Nov 18, 2018 · 3668f4d · 3668f4d
1 parent a8b27ec
commit 3668f4d
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 18 deletions.
diff --git a/src/predictor/KerasPredictor.js b/src/predictor/KerasPredictor.js
@@ -43,13 +43,19 @@ KerasPredictor.prototype = Object.assign( Object.create( Predictor.prototype ),
 
 	predict: function( data ) {
 
-		// Create input tensor for prediction.
+		let predictor = this;
 
-		let inputTensor = this.createInputTensor( data );
+		let predictResult = tf.tidy( () => {
 
-		// Get prediction result from loaded model.
+			// Create input tensor for prediction.
 
-		let predictResult = this.model.resource.predict( inputTensor );
+			let inputTensor = predictor.createInputTensor( data );
+
+			// Get prediction result from loaded model.
+
+			return predictor.model.resource.predict( inputTensor );
+
+		} );
 
 		return predictResult;
 

diff --git a/src/predictor/TfPredictor.js b/src/predictor/TfPredictor.js
@@ -52,25 +52,33 @@ TfPredictor.prototype = Object.assign( Object.create( Predictor.prototype ), {
 
 	predict: function( data ) {
 
-		// Create input tensor for prediction.
+		let predictor = this;
 
-		let inputTensor = this.createInputTensor( data );
+		let predictResult = tf.tidy( () => {
 
-		let predictResult;
+			// Create input tensor for prediction.
 
-		if ( this.outputsName !== undefined ) {
+			let inputTensor = predictor.createInputTensor( data );
 
-			// If has outputsName, use execute to get prediction result.
+			let predictResult;
 
-			predictResult = this.model.resource.execute( inputTensor, this.outputsName );
+			if ( this.outputsName !== undefined ) {
 
-		} else {
+				// If has outputsName, use execute to get prediction result.
 
-			// If outputsName is undefined, use predict to get prediction result.
+				predictResult = predictor.model.resource.execute( inputTensor, this.outputsName );
 
-			predictResult = this.model.resource.predict( inputTensor );
+			} else {
 
-		}
+				// If outputsName is undefined, use predict to get prediction result.
+
+				predictResult = predictor.model.resource.predict( inputTensor );
+
+			}
+
+			return predictResult;
+
+		} );
 
 		return predictResult;
 

diff --git a/src/predictor/TfjsPredictor.js b/src/predictor/TfjsPredictor.js
@@ -43,13 +43,19 @@ TfjsPredictor.prototype = Object.assign( Object.create( Predictor.prototype ), {
 
 	predict: function( data ) {
 
-		// Create input tensor for prediction.
+		let predictor = this;
 
-		let inputTensor = this.createInputTensor( data );
+		let predictResult = tf.tidy( () => {
 
-		// Get prediction result from loaded model.
+			// Create input tensor for prediction.
 
-		let predictResult = this.model.resource.predict( inputTensor );
+			let inputTensor = predictor.createInputTensor( data );
+
+			// Get prediction result from loaded model.
+
+			return predictor.model.resource.predict( inputTensor );
+
+		} );
 
 		return predictResult;