tensorflow: [Java API] Tensor.create() slow for large arrays

Ok, yeah it looks like the way to encode a TF string in a buffer is way more complex that I first thought 😕 Basically, it follows as structure like this (for String array [A, B, C]):

[offset A][offset B][offset C][length A][data A][length B][data B][length C][data C]

After a few tests, I think the following code should do the job, I’m curious that you give it a try in your projects and if 1) it works and 2) it is faster, I can push it as a utility in the repo. Also note that there is no padding in the following algorithm, could be added if needed:

  private static ByteBuffer stringArrayToBuffer(String[] values) throws IOException {
    long offsets[] = new long[values.length];
    byte[][] data = new byte[values.length][];
    int dataSize = 0;

    // Convert strings to encoded bytes and calculate required data size, including a varint for each of them
    for (int i = 0; i < values.length; ++i) {
      byte[] byteValue = values[i].getBytes("UTF-8");
      data[i] = byteValue;
      int length = byteValue.length + varintLength(byteValue.length);
      dataSize += length;
      if (i < values.length - 1) {
        offsets[i + 1] = offsets[i] + length;
      }
    }

    // Important: buffer must follow native byte order
    ByteBuffer buffer = ByteBuffer.allocate(dataSize + (offsets.length * 8)).order(ByteOrder.nativeOrder());

    // First, write offsets to each elements in the buffer
    for (int i = 0; i < offsets.length; ++i) {
      buffer.putLong(offsets[i]);
    }
    
    // Second, write strings bytes, each preceded by its length encoded as a varint
    for (int i = 0; i < data.length; ++i) {
      encodeVarint(buffer, data[i].length);
      buffer.put(data[i]);
    }

    return (ByteBuffer)buffer.rewind();
  }
  
  private static void encodeVarint(ByteBuffer buffer, int value) {
      int v = value;
      while (v >= 0x80) {
        buffer.put((byte)((v & 0x7F) | 0x80));
        v >>= 7;
      }
      buffer.put((byte)v);
  }

  private static int varintLength(int length) {
    int len = 1;
    while (length >= 0x80) {
      length >>= 7;
      ++len;
    }
    return len;
  }

Thanks for the detailed description and the sample code, it is very much appreciated!

The create(Object) method call involves use of reflection to determine the shape and copy things over one array at a time, so it is pretty slow, especially as you add dimensions. The create(shape, FloatBuffer) method would be an order-of-magnitude faster. For example:

public void test() {
    Random r = new Random();
    int imageSize = 224 * 224 * 3;
    int batch = 128;
    long[] shape = new long[] {batch, imageSize};
    FloatBuffer buf = FloatBuffer.allocate(imageSize * batch);
    for (int i = 0; i < imageSize * batch; ++i) {
      buf.put(r.nextFloat());
    }
    buf.flip();

    long start = System.nanoTime();
    Tensor.create(shape, buf);
    long end = System.nanoTime();
    System.out.println("Took: " + (end - start));
}

This is still slower than I’d want it to be, have to dig into that, but hopefully it is enough to satisfy your current needs (and session execution should be faster than Python).

Regarding your other question: Yes, using feeds is slower than getting input from queues. While I believe we do have the primitives to enable use of queues from any language, it is admittedly not too easy (as you have to figure out what is being done in python - e.g., start threads that run the enqueue op - and duplicate that). Note that, as per the proposal in #7951 - investing into queues in other languages might not be worth while at this stage.

Do let me know if using the FloatBuffer suffices for now (and I can close this issue, while we look into general performance improvements for the Java API).

By the way, all this is currently being addressed in the new official Java repository for TensorFlow, there is a important project called “Tensor NIO” that I’m working on which will allow users to directly access tensor memory from Java and to read/write its data in an N-dimensional space. That will include what is found in the Buffers class of the previous example.

If anyone is interested to know more about it, just let me know.

@maziyarpanahi : sure, it’s now here: https://github.com/karllessard/tensorflow-models/tree/master/samples/languages/java/suite/src/main/java/org/tensorflow/model/sample/buffer

@asimshankar this is really great stuff, the timings improved exactly by an order of magnitude as you said, thanks! A batch of 128 takes around 0.2s and the batch of 32 takes around 0.04s. Still I think a bit slower than what Python does but definitely I can work with that, thanks!

Maybe you could add a note in the Java doc about this? I was suspecting the reflection to take a long time in Tensorflow.create(Object) but wasn’t sure if that was the main cause.