concat score tensors before getting their data, apparently it's faster this way

justadudewhohacks · justadudewhohacks · commit e7d1d0437ed3 · 2018-07-12T17:36:22.000+02:00
diff --git a/src/mtcnn/stage2.ts b/src/mtcnn/stage2.ts
@@ -1,3 +1,4 @@
+import { tf } from '..';
 import { BoundingBox } from './BoundingBox';
 import { extractImagePatches } from './extractImagePatches';
 import { nms } from './nms';
@@ -26,8 +27,12 @@ export async function stage2(
   )
   stats.stage2_rnet = Date.now() - ts
 
-  const scoreDatas = await Promise.all(rnetOuts.map(out => out.scores.data()))
-  const scores = scoreDatas.map(arr => Array.from(arr)).reduce((all, arr) => all.concat(arr))
+  const scoresTensor = rnetOuts.length > 1
+    ? tf.concat(rnetOuts.map(out => out.scores))
+    : rnetOuts[0].scores
+  const scores = Array.from(await scoresTensor.data())
+  scoresTensor.dispose()
+
   const indices = scores
     .map((score, idx) => ({ score, idx }))
     .filter(c => c.score > scoreThreshold)
diff --git a/src/mtcnn/stage3.ts b/src/mtcnn/stage3.ts
@@ -4,6 +4,7 @@ import { extractImagePatches } from './extractImagePatches';
 import { nms } from './nms';
 import { ONet } from './ONet';
 import { ONetParams } from './types';
+import { tf } from '..';
 
 export async function stage3(
   img: HTMLCanvasElement,
@@ -27,8 +28,12 @@ export async function stage3(
   )
   stats.stage3_onet = Date.now() - ts
 
-  const scoreDatas = await Promise.all(onetOuts.map(out => out.scores.data()))
-  const scores = scoreDatas.map(arr => Array.from(arr)).reduce((all, arr) => all.concat(arr))
+  const scoresTensor = onetOuts.length > 1
+    ? tf.concat(onetOuts.map(out => out.scores))
+    : onetOuts[0].scores
+  const scores = Array.from(await scoresTensor.data())
+  scoresTensor.dispose()
+
   const indices = scores
     .map((score, idx) => ({ score, idx }))
     .filter(c => c.score > scoreThreshold)
