diff options
author | 2018-03-30 11:21:31 -0700 | |
---|---|---|
committer | 2018-03-30 11:24:25 -0700 | |
commit | ff451cae700d0f0f17ff9d2dde32299344b21fb1 (patch) | |
tree | af2f428211f3f896fc4116ccd792d0e55bfed1eb /tensorflow/contrib/lite/java | |
parent | c4acdccbb7284c6a63e6824a7ee45ce7a86606b9 (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.java | 13 | ||||
-rw-r--r-- | tensorflow/contrib/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java | 12 |
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); } /** |