test new builders

bioimage-io · Oct 3, 2024 · be23fbf · be23fbf
1 parent 6516637
commit be23fbf
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 7 deletions.
diff --git a/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api030/shm/ShmBuilder.java b/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api030/shm/ShmBuilder.java
@@ -28,6 +28,7 @@
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 
+import org.tensorflow.ndarray.buffer.ByteDataBuffer;
 import org.tensorflow.types.TFloat32;
 import org.tensorflow.types.TFloat64;
 import org.tensorflow.types.TInt32;
@@ -101,7 +102,19 @@ private static void buildFromTensorUByte(TUint8 tensor, String memoryName) throw
 					+ " is too big. Max number of elements per ubyte output tensor supported: " + Integer.MAX_VALUE / 1);
         SharedMemoryArray shma = SharedMemoryArray.readOrCreate(memoryName, arrayShape, new UnsignedByteType(), false, true);
         ByteBuffer buff = shma.getDataBufferNoHeader();
-        tensor.asRawTensor().data().read(buff.array(), 0, buff.capacity());
+        long tt = System.currentTimeMillis();
+        ByteDataBuffer tensorData = tensor.asRawTensor().data();
+        for (int i = 0; i < buff.capacity(); i ++) {
+        	buff.put(tensorData.getByte(i));
+        }
+        System.out.println("TIME 1: " + (System.currentTimeMillis() - tt) / 1000);
+        buff.rewind();
+        tt = System.currentTimeMillis();
+        byte[] flat = new byte[buff.capacity()];
+        ByteBuffer buff2 = ByteBuffer.wrap(flat);
+        tensor.asRawTensor().data().read(flat, 0, buff.capacity());
+        shma.setBuffer(buff2);
+        System.out.println("TIME 2: " + (System.currentTimeMillis() - tt) / 1000);
         if (PlatformDetection.isWindows()) shma.close();
     }
 
@@ -127,7 +140,19 @@ private static void buildFromTensorFloat(TFloat32 tensor, String memoryName) thr
 
         SharedMemoryArray shma = SharedMemoryArray.readOrCreate(memoryName, arrayShape, new FloatType(), false, true);
         ByteBuffer buff = shma.getDataBufferNoHeader();
-        tensor.asRawTensor().data().read(buff.array(), 0, buff.capacity());
+        long tt = System.currentTimeMillis();
+        ByteDataBuffer tensorData = tensor.asRawTensor().data();
+        for (int i = 0; i < buff.capacity(); i ++) {
+        	buff.put(tensorData.getByte(i));
+        }
+        System.out.println("TIME 1: " + (System.currentTimeMillis() - tt) / 1000);
+        buff.rewind();
+        tt = System.currentTimeMillis();
+        byte[] flat = new byte[buff.capacity()];
+        ByteBuffer buff2 = ByteBuffer.wrap(flat);
+        tensor.asRawTensor().data().read(flat, 0, buff.capacity());
+        shma.setBuffer(buff2);
+        System.out.println("TIME 2: " + (System.currentTimeMillis() - tt) / 1000);
         if (PlatformDetection.isWindows()) shma.close();
     }
 

diff --git a/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api030/shm/TensorBuilder.java b/src/main/java/io/bioimage/modelrunner/tensorflow/v2/api030/shm/TensorBuilder.java
@@ -97,7 +97,10 @@ private static TUint8 buildUByte(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		ByteDataBuffer dataBuffer = RawDataBufferFactory.create(buff.array(), false);
+		byte[] flat = new byte[buff.capacity()];
+		buff.get(flat);
+		buff.rewind();
+		ByteDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		TUint8 ndarray = Tensor.of(TUint8.class, Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
@@ -112,7 +115,10 @@ private static TInt32 buildInt(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		IntDataBuffer dataBuffer = RawDataBufferFactory.create(buff.asIntBuffer().array(), false);
+		int[] flat = new int[buff.capacity() / 4];
+		buff.asIntBuffer().get(flat);
+		buff.rewind();
+		IntDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		TInt32 ndarray = TInt32.tensorOf(Shape.of(ogShape),
 			dataBuffer);
 		return ndarray;
@@ -128,7 +134,10 @@ private static TInt64 buildLong(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		LongDataBuffer dataBuffer = RawDataBufferFactory.create(buff.asLongBuffer().array(), false);
+		long[] flat = new long[buff.capacity() / 8];
+		buff.asLongBuffer().get(flat);
+		buff.rewind();
+		LongDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		TInt64 ndarray = TInt64.tensorOf(Shape.of(ogShape),
 			dataBuffer);
 		return ndarray;
@@ -144,7 +153,10 @@ private static TFloat32 buildFloat(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		FloatDataBuffer dataBuffer = RawDataBufferFactory.create(buff.asFloatBuffer().array(), false);
+		float[] flat = new float[buff.capacity() / 4];
+		buff.asFloatBuffer().get(flat);
+		buff.rewind();
+		FloatDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		TFloat32 ndarray = TFloat32.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}
@@ -159,7 +171,10 @@ private static TFloat64 buildDouble(SharedMemoryArray tensor)
 		if (!tensor.isNumpyFormat())
 			throw new IllegalArgumentException("Shared memory arrays must be saved in numpy format.");
 		ByteBuffer buff = tensor.getDataBufferNoHeader();
-		DoubleDataBuffer dataBuffer = RawDataBufferFactory.create(buff.asDoubleBuffer().array(), false);
+		double[] flat = new double[buff.capacity() / 8];
+		buff.asDoubleBuffer().get(flat);
+		buff.rewind();
+		DoubleDataBuffer dataBuffer = RawDataBufferFactory.create(flat, false);
 		TFloat64 ndarray = TFloat64.tensorOf(Shape.of(ogShape), dataBuffer);
 		return ndarray;
 	}