How to allocate multi dimensional outputbuffer to feed Android Tflite's interpreter.run()?

Question

I try to run a version posenet (which is a CNN) on an android app with tflite.

The app is based on the GPU delegate demo:

(1) https://medium.com/tensorflow/tensorflow-lite-now-faster-with-mobile-gpus-developer-preview-e15797e6dee7

(2) https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/java/demo

Posenet takes an Image as Input and computes as output multiple arrays of the shape: 1x14x14x17, 1x14x14x34, 1x14x14x32, 1x14x14x32

I know how to allocate a bytebuffer for the image, so thats not the problem.

But how do i allocate a buffer for that output to be able to successfully feed the input and output buffer to the interpreter like:

import org.tensorflow.lite.Interpreter;
Interpreter tflite;

ByteBuffer input = null;
input = ByteBuffer.allocateDirect(...);
output = ?

tflite.run(input,output);

I tried something like this for the float version:

float[][][][] output = null;
output = new float[1*14*14*17][1*14*14*34][1*14*14*32][1*14*14*32];

but this leads to a memory oom. So how do I correctly allocate a buffe for the output with an array with the right dimensions. (I am not so used to java, more to python)

EDIT:

I want four nested 4-dimensional arrays. Because thats whats coming out of the model. something like this:

float[][][][] out1 = new float[1][14][14][17];
float[][][][] out2 = new float[1][14][14][34];
float[][][][] out3 = new float[1][14][14][32];
float[][][][] out4 = new float[1][14][14][32];
float[] output = new float[out1, out2, out3, out4];

But this does not work

Answer 1

@Chris623's answer still would not work for use in tflite. Using interpreter.run() would not work in the case where you need to allocate memory for a multi-dimensional output array.

You would rather need to use interpreter.runForMultipleInputsOutputs(). Sample code is shown below for your 4 arrays.

float[][][][] out1 = new float[1][14][14][17];
float[][][][] out2 = new float[1][14][14][34];
float[][][][] out3 = new float[1][14][14][32];
float[][][][] out4 = new float[1][14][14][32];

Map<Integer, Object> outputs = new HashMap<>();
outputs.put(0, out1)
outputs.put(1, out2)
outputs.put(2, out3)
outputs.put(3, out4)

interpreter.runForMultipleInputsOutputs(inputs, outputs)

This should help you map your output arrays for further processing.