aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/lite/java
diff options
context:
space:
mode:
authorGravatar Shashi Shekhar <shashishekhar@google.com>2018-03-30 11:21:31 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-03-30 11:24:25 -0700
commitff451cae700d0f0f17ff9d2dde32299344b21fb1 (patch)
treeaf2f428211f3f896fc4116ccd792d0e55bfed1eb /tensorflow/contrib/lite/java
parentc4acdccbb7284c6a63e6824a7ee45ce7a86606b9 (diff)
Internal change.
PiperOrigin-RevId: 191090993
Diffstat (limited to 'tensorflow/contrib/lite/java')
-rw-r--r--tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java13
-rw-r--r--tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java12
2 files changed, 19 insertions, 6 deletions
diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
index 14f461f5f9..a33959dca4 100644
--- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
+++ b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java
@@ -68,6 +68,19 @@ public final class Interpreter implements AutoCloseable {
}
/**
+ * Initializes a {@code Interpreter} and specifies the number of threads used for inference.
+ *
+ * @param modelFile: a file of a pre-trained TF Lite model
+ * @param numThreads: number of threads to use for inference
+ */
+ public Interpreter(@NonNull File modelFile, int numThreads) {
+ if (modelFile == null) {
+ return;
+ }
+ wrapper = new NativeInterpreterWrapper(modelFile.getAbsolutePath(), numThreads);
+ }
+
+ /**
* Initializes a {@code Interpreter} with a {@code MappedByteBuffer} to the model file.
*
* <p>The {@code MappedByteBuffer} should remain unchanged after the construction of a {@code
diff --git a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
index dbf8f8f7cc..fc8187acfe 100644
--- a/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
+++ b/tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
@@ -32,9 +32,13 @@ import java.util.Map;
final class NativeInterpreterWrapper implements AutoCloseable {
NativeInterpreterWrapper(String modelPath) {
+ this(modelPath, /* numThreads= */ -1);
+ }
+
+ NativeInterpreterWrapper(String modelPath, int numThreads) {
errorHandle = createErrorReporter(ERROR_BUFFER_SIZE);
modelHandle = createModel(modelPath, errorHandle);
- interpreterHandle = createInterpreter(modelHandle, errorHandle, /* numThreads= */ -1);
+ interpreterHandle = createInterpreter(modelHandle, errorHandle, numThreads);
isMemoryAllocated = true;
}
@@ -44,11 +48,7 @@ final class NativeInterpreterWrapper implements AutoCloseable {
* NativeInterpreterWrapper}.
*/
NativeInterpreterWrapper(MappedByteBuffer mappedByteBuffer) {
- modelByteBuffer = mappedByteBuffer;
- errorHandle = createErrorReporter(ERROR_BUFFER_SIZE);
- modelHandle = createModelWithBuffer(modelByteBuffer, errorHandle);
- interpreterHandle = createInterpreter(modelHandle, errorHandle, /* numThreads= */ -1);
- isMemoryAllocated = true;
+ this(mappedByteBuffer, /* numThreads= */ -1);
}
/**