added some tests to indentify memory leaks in NetInput and face landmark net +…

src/faceLandmarkNet/FaceLandmarkNet.ts

@@ -108,7 +108,9 @@ export class FaceLandmarkNet {
   public async detectLandmarks(input: TNetInput): Promise<FaceLandmarks | FaceLandmarks[]> {
     const netInput = await toNetInput(input, true)
 
-    const landmarkTensors = tf.unstack(this.forwardInput(netInput))
+    const landmarkTensors = tf.tidy(
+      () => tf.unstack(this.forwardInput(netInput))
+    )
 
     const landmarksForBatch = await Promise.all(landmarkTensors.map(
       async (landmarkTensor, batchIdx) => {

test/tests/NetInput.test.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { NetInput } from '../../src/NetInput';
+import { bufferToImage, createCanvasFromMedia } from '../../src/utils';
+import { expectAllTensorsReleased, tensor3D } from '../utils';
+
+
+describe('NetInput', () => {
+
+  let imgEl: HTMLImageElement, canvasEl: HTMLCanvasElement
+
+  beforeAll(async () => {
+    const img = await (await fetch('base/test/images/face1.png')).blob()
+    imgEl = await bufferToImage(img)
+    canvasEl = createCanvasFromMedia(imgEl)
+  })
+
+  describe('no memory leaks', () => {
+
+    describe('constructor', () => {
+
+      it('single image element', async () => {
+        await expectAllTensorsReleased(() => {
+          const net = new NetInput([imgEl])
+          net.dispose()
+        })
+      })
+
+      it('multiple image elements', async () => {
+        await expectAllTensorsReleased(() => {
+          const net = new NetInput([imgEl, imgEl, imgEl])
+          net.dispose()
+        })
+      })
+
+      it('single tf.Tensor3D', async () => {
+        const tensor = tensor3D()
+
+        await expectAllTensorsReleased(() => {
+          const net = new NetInput([tensor])
+          net.dispose()
+        })
+
+        tensor.dispose()
+      })
+
+      it('multiple tf.Tensor3Ds', async () => {
+        const tensors = [tensor3D(), tensor3D(), tensor3D()]
+
+        await expectAllTensorsReleased(() => {
+          const net = new NetInput(tensors)
+          net.dispose()
+        })
+
+        tensors.forEach(t => t.dispose())
+      })
+    })
+
+    describe('toBatchTensor', () => {
+
+      it('single image element', async () => {
+        await expectAllTensorsReleased(() => {
+          const net = new NetInput([imgEl])
+          const batchTensor = net.toBatchTensor(100, false)
+          net.dispose()
+          batchTensor.dispose()
+        })
+      })
+
+      it('multiple image elements', async () => {
+        await expectAllTensorsReleased(() => {
+          const net = new NetInput([imgEl, imgEl, imgEl])
+          const batchTensor = net.toBatchTensor(100, false)
+          net.dispose()
+          batchTensor.dispose()
+        })
+      })
+
+      it('managed, single image element', async () => {
+        await expectAllTensorsReleased(() => {
+          const net = (new NetInput([imgEl])).managed()
+          const batchTensor = net.toBatchTensor(100, false)
+          batchTensor.dispose()
+        })
+      })
+
+      it('managed, multiple image elements', async () => {
+        await expectAllTensorsReleased(() => {
+          const net = (new NetInput([imgEl, imgEl, imgEl])).managed()
+          const batchTensor = net.toBatchTensor(100, false)
+          batchTensor.dispose()
+        })
+      })
+
+    })
+
+  })
+
+})

test/tests/e2e/faceLandmarkNet.test.ts

@@ -5,7 +5,8 @@ import { isTensor3D } from '../../../src/commons/isTensor';
 import { FaceLandmarks } from '../../../src/faceLandmarkNet/FaceLandmarks';
 import { Point } from '../../../src/Point';
 import { Dimensions, TMediaElement } from '../../../src/types';
-import { expectMaxDelta } from '../../utils';
+import { expectMaxDelta, expectAllTensorsReleased } from '../../utils';
+import { NetInput } from '../../../src/NetInput';
 
 function getInputDims (input: tf.Tensor | TMediaElement): Dimensions {
   if (input instanceof tf.Tensor) {
@@ -115,7 +116,6 @@ describe('faceLandmarkNet', () => {
 
   })
 
-
   describe('batch inputs', () => {
 
     let faceLandmarkNet: faceapi.FaceLandmarkNet
@@ -134,6 +134,7 @@ describe('faceLandmarkNet', () => {
         faceLandmarkPositions2,
         faceLandmarkPositionsRect
       ]
+
       const results = await faceLandmarkNet.detectLandmarks(inputs) as FaceLandmarks[]
       expect(Array.isArray(results)).toBe(true)
       expect(results.length).toEqual(3)
@@ -158,6 +159,7 @@ describe('faceLandmarkNet', () => {
         faceLandmarkPositions2,
         faceLandmarkPositionsRect
       ]
+
       const results = await faceLandmarkNet.detectLandmarks(inputs) as FaceLandmarks[]
       expect(Array.isArray(results)).toBe(true)
       expect(results.length).toEqual(3)
@@ -182,6 +184,7 @@ describe('faceLandmarkNet', () => {
         faceLandmarkPositions2,
         faceLandmarkPositionsRect
       ]
+
       const results = await faceLandmarkNet.detectLandmarks(tf.stack(inputs) as tf.Tensor4D) as FaceLandmarks[]
       expect(Array.isArray(results)).toBe(true)
       expect(results.length).toEqual(2)
@@ -207,7 +210,6 @@ describe('faceLandmarkNet', () => {
         faceLandmarkPositionsRect
       ]
 
-
       const results = await faceLandmarkNet.detectLandmarks(inputs) as FaceLandmarks[]
       expect(Array.isArray(results)).toBe(true)
       expect(results.length).toEqual(3)
@@ -226,5 +228,51 @@ describe('faceLandmarkNet', () => {
 
   })
 
+  describe('no memory leaks', () => {
+
+    let faceLandmarkNet: faceapi.FaceLandmarkNet
+
+    beforeAll(async () => {
+      faceLandmarkNet = new faceapi.FaceLandmarkNet()
+      await faceLandmarkNet.load('base/weights')
+    })
+
+    describe('forwardInput', () => {
+
+      it('single image element', async () => {
+        await expectAllTensorsReleased(async () => {
+          const netInput = (new NetInput([imgEl1])).managed()
+          const outTensor = await faceLandmarkNet.forwardInput(netInput)
+          outTensor.dispose()
+        })
+      })
+
+      it('multiple image elements', async () => {
+        await expectAllTensorsReleased(async () => {
+          const netInput = (new NetInput([imgEl1, imgEl1, imgEl1])).managed()
+          const outTensor = await faceLandmarkNet.forwardInput(netInput)
+          outTensor.dispose()
+        })
+      })
+
+    })
+
+    describe('detectLandmarks', () => {
+
+      it('single image element', async () => {
+        await expectAllTensorsReleased(async () => {
+          await faceLandmarkNet.detectLandmarks(imgEl1)
+        })
+      })
+
+      it('multiple image elements', async () => {
+        await expectAllTensorsReleased(async () => {
+          await faceLandmarkNet.detectLandmarks([imgEl1, imgEl1, imgEl1])
+        })
+      })
+
+    })
+  })
+
 })
 

test/tests/toNetInput.test.ts

 import { NetInput } from '../../src/NetInput';
 import { toNetInput } from '../../src/toNetInput';
 import { bufferToImage, createCanvasFromMedia } from '../../src/utils';
-
-async function createFakeHTMLVideoElement() {
-  const videoEl = document.createElement('video')
-  videoEl.muted = true
-  videoEl.src = 'base/test/media/video.mp4'
-  await videoEl.pause()
-  await videoEl.play()
-  return videoEl
-}
+import { createFakeHTMLVideoElement } from '../utils';
 
 describe('toNetInput', () => {
 

test/utils.ts

+import * as tf from '@tensorflow/tfjs-core';
+
 export function zeros(length: number): Float32Array {
   return new Float32Array(length)
 }
@@ -9,3 +11,22 @@ export function ones(length: number): Float32Array {
 export function expectMaxDelta(val1: number, val2: number, maxDelta: number) {
   expect(Math.abs(val1 - val2)).toBeLessThan(maxDelta)
 }
+
+export async function createFakeHTMLVideoElement() {
+  const videoEl = document.createElement('video')
+  videoEl.muted = true
+  videoEl.src = 'base/test/media/video.mp4'
+  await videoEl.pause()
+  await videoEl.play()
+  return videoEl
+}
+
+export async function expectAllTensorsReleased(fn: () => any) {
+  const numTensorsBefore = tf.memory().numTensors
+  await fn()
+  expect(tf.memory().numTensors - numTensorsBefore).toEqual(0)
+}
+
+export function tensor3D() {
+  return tf.tensor3d([[[0]]])
+}
\ No newline at end of file