diff --git a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java index 6aeb06355b4..efcdc0e4c65 100644 --- a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java +++ b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/Interpreter.java @@ -314,6 +314,32 @@ public final class Interpreter implements AutoCloseable { wrapper.run(inputs, outputs); } + /** + * Expicitly updates allocations for all tensors, if necessary. + * + *
This will propagate shapes and memory allocations for all dependent tensors using the input + * tensor shape(s) as given. + * + *
Note: This call is *purely optional*. Tensor allocation will occur automatically during + * execution if any input tensors have been resized. This call is most useful in determining the + * shapes for any output tensors before executing the graph, e.g., + *
{@code + * interpreter.resizeInput(0, new int[]{1, 4, 4, 3})); + * interpreter.allocateTensors(); + * FloatBuffer input = FloatBuffer.allocate(interpreter.getInputTensor(0),numElements()); + * // Populate inputs... + * FloatBuffer output = FloatBuffer.allocate(interpreter.getOutputTensor(0).numElements()); + * interpreter.run(input, output) + * // Process outputs... + * }+ * + * @throws IllegalStateException if the graph's tensors could not be successfully allocated. + */ + public void allocateTensors() { + checkNotClosed(); + wrapper.allocateTensors(); + } + /** * Resizes idx-th input of the native model to the given dims. * @@ -373,6 +399,13 @@ public final class Interpreter implements AutoCloseable { /** * Gets the Tensor associated with the provdied output index. * + *
Note: Output tensor details (e.g., shape) may not be fully populated until after inference
+ * is executed. If you need updated details *before* running inference (e.g., after resizing an
+ * input tensor, which may invalidate output tensor shapes), use {@link #allocateTensors()} to
+ * explicitly trigger allocation and shape propagation. Note that, for graphs with output shapes
+ * that are dependent on input *values*, the output shape may not be fully determined until
+ * running inference.
+ *
* @throws IllegalArgumentException if {@code outputIndex} is negtive or is not smaller than the
* number of model outputs.
*/
diff --git a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
index ca21ec5c7ea..73fe506f131 100644
--- a/tensorflow/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
+++ b/tensorflow/lite/java/src/main/java/org/tensorflow/lite/NativeInterpreterWrapper.java
@@ -175,6 +175,8 @@ final class NativeInterpreterWrapper implements AutoCloseable {
/** Resizes dimensions of a specific input. */
void resizeInput(int idx, int[] dims) {
if (resizeInput(interpreterHandle, errorHandle, idx, dims)) {
+ // Tensor allocation is deferred until either an explicit `allocateTensors()` call or
+ // `invoke()` avoiding redundant allocations if multiple tensors are simultaneosly resized.
isMemoryAllocated = false;
if (inputTensors[idx] != null) {
inputTensors[idx].refreshShape();
@@ -185,6 +187,23 @@ final class NativeInterpreterWrapper implements AutoCloseable {
private static native boolean resizeInput(
long interpreterHandle, long errorHandle, int inputIdx, int[] dims);
+ /** Triggers explicit allocation of tensors. */
+ void allocateTensors() {
+ if (isMemoryAllocated) {
+ return;
+ }
+
+ isMemoryAllocated = true;
+ allocateTensors(interpreterHandle, errorHandle);
+ for (int i = 0; i < outputTensors.length; ++i) {
+ if (outputTensors[i] != null) {
+ outputTensors[i].refreshShape();
+ }
+ }
+ }
+
+ private static native long allocateTensors(long interpreterHandle, long errorHandle);
+
void setUseNNAPI(boolean useNNAPI) {
useNNAPI(interpreterHandle, useNNAPI);
}
@@ -385,8 +404,6 @@ final class NativeInterpreterWrapper implements AutoCloseable {
// List of owned delegates that must be closed when the interpreter is closed.
private final List