build param mappings for remaining nets (detection, recognition) to make them disposable

src/Rect.ts

@@ -32,6 +32,11 @@ export class Rect implements IRect {
     return new Rect(x, y, width, height)
   }
 
+  public pad(padX: number, padY: number): Rect {
+    let { x, y, width, height } = this
+    return new Rect(x - (padX / 2), y - (padY / 2), width + padX, height + padY)
+  }
+
   public floor(): Rect {
     return new Rect(
       Math.floor(this.x),

src/commons/NeuralNetwork.ts

@@ -7,6 +7,8 @@ export class NeuralNetwork<TNetParams> {
   protected _params: TNetParams | undefined = undefined
   protected _paramMappings: ParamMapping[] = []
 
+  constructor(private _name: string) {}
+
   public get params(): TNetParams | undefined {
     return this._params
   }
@@ -53,11 +55,44 @@ export class NeuralNetwork<TNetParams> {
     })
   }
 
-  public dispose() {
-    this.getParamList().forEach(param => param.tensor.dispose())
+  public dispose(throwOnRedispose: boolean = true) {
+    this.getParamList().forEach(param => {
+      if (throwOnRedispose && param.tensor.isDisposed) {
+        throw new Error(`param tensor has already been disposed for path ${param.path}`)
+      }
+      param.tensor.dispose()
+    })
     this._params = undefined
   }
 
+  public async load(weightsOrUrl: Float32Array | string | undefined): Promise<void> {
+    if (weightsOrUrl instanceof Float32Array) {
+      this.extractWeights(weightsOrUrl)
+      return
+    }
+
+    if (weightsOrUrl && typeof weightsOrUrl !== 'string') {
+      throw new Error(`${this._name}.load - expected model uri, or weights as Float32Array`)
+    }
+    const {
+      paramMappings,
+      params
+    } = await this.loadQuantizedParams(weightsOrUrl)
+
+    this._paramMappings = paramMappings
+    this._params = params
+  }
+
+  public extractWeights(weights: Float32Array) {
+    const {
+      paramMappings,
+      params
+    } = this.extractParams(weights)
+
+    this._paramMappings = paramMappings
+    this._params = params
+  }
+
   private traversePropertyPath(paramPath: string) {
     if (!this.params) {
       throw new Error(`traversePropertyPath - model has no loaded params`)
@@ -78,4 +113,12 @@ export class NeuralNetwork<TNetParams> {
 
     return { obj, objProp }
   }
+
+  protected loadQuantizedParams(_: any): Promise<{ params: TNetParams, paramMappings: ParamMapping[] }> {
+    throw new Error(`${this._name}.loadQuantizedParams - not implemented`)
+  }
+
+  protected extractParams(_: any): { params: TNetParams, paramMappings: ParamMapping[] } {
+    throw new Error(`${this._name}.extractParams - not implemented`)
+  }
 }
\ No newline at end of file

src/commons/disposeUnusedWeightTensors.ts

+import { ParamMapping } from './types';
+
+export function disposeUnusedWeightTensors(weightMap: any, paramMappings: ParamMapping[]) {
+  Object.keys(weightMap).forEach(path => {
+    if (!paramMappings.some(pm => pm.originalPath === path)) {
+      weightMap[path].dispose()
+    }
+  })
+}

src/commons/extractConvParamsFactory.ts

-import * as tf from '@tensorflow/tfjs-core';
-
-import { ConvParams, ExtractWeightsFunction, ParamMapping } from './types';
-
-export function extractConvParamsFactory(extractWeights: ExtractWeightsFunction, paramMappings: ParamMapping[]) {
-  return function (
-    channelsIn: number,
-    channelsOut: number,
-    filterSize: number,
-    mappedPrefix: string
-  ): ConvParams {
-    const filters = tf.tensor4d(
-      extractWeights(channelsIn * channelsOut * filterSize * filterSize),
-      [filterSize, filterSize, channelsIn, channelsOut]
-    )
-    const bias = tf.tensor1d(extractWeights(channelsOut))
-
-    paramMappings.push(
-      { paramPath: `${mappedPrefix}/filters` },
-      { paramPath: `${mappedPrefix}/bias` }
-    )
-
-    return {
-      filters,
-      bias
-    }
-  }
-}
\ No newline at end of file

src/commons/extractWeightEntry.ts

-import { isTensor } from './isTensor';
-
-export function extractWeightEntry(weightMap: any, path: string, paramRank: number) {
-  const tensor = weightMap[path]
-
-  if (!isTensor(tensor, paramRank)) {
-    throw new Error(`expected weightMap[${path}] to be a Tensor${paramRank}D, instead have ${tensor}`)
-  }
-
-  return { path, tensor }
-}
\ No newline at end of file

src/commons/extractWeightEntryFactory.ts

+import { isTensor } from './isTensor';
+import { ParamMapping } from './types';
+
+export function extractWeightEntryFactory(weightMap: any, paramMappings: ParamMapping[]) {
+
+  return function<T> (originalPath: string, paramRank: number, mappedPath?: string): T {
+    const tensor = weightMap[originalPath]
+
+    if (!isTensor(tensor, paramRank)) {
+      throw new Error(`expected weightMap[${originalPath}] to be a Tensor${paramRank}D, instead have ${tensor}`)
+    }
+
+    paramMappings.push(
+      { originalPath, paramPath: mappedPath || originalPath }
+    )
+
+    return tensor
+  }
+
+}

src/faceDetectionNet/FaceDetectionNet.ts

 import * as tf from '@tensorflow/tfjs-core';
 
+import { NeuralNetwork } from '../commons/NeuralNetwork';
 import { NetInput } from '../NetInput';
 import { Rect } from '../Rect';
 import { toNetInput } from '../toNetInput';
@@ -13,28 +14,17 @@ import { outputLayer } from './outputLayer';
 import { predictionLayer } from './predictionLayer';
 import { NetParams } from './types';
 
-export class FaceDetectionNet {
+export class FaceDetectionNet extends NeuralNetwork<NetParams> {
 
-  private _params: NetParams
-
-  public async load(weightsOrUrl?: Float32Array | string): Promise<void> {
-    if (weightsOrUrl instanceof Float32Array) {
-      this.extractWeights(weightsOrUrl)
-      return
-    }
-
-    if (weightsOrUrl && typeof weightsOrUrl !== 'string') {
-      throw new Error('FaceDetectionNet.load - expected model uri, or weights as Float32Array')
-    }
-    this._params = await loadQuantizedParams(weightsOrUrl)
-  }
-
-  public extractWeights(weights: Float32Array) {
-    this._params = extractParams(weights)
+  constructor() {
+    super('FaceDetectionNet')
   }
 
   public forwardInput(input: NetInput) {
-    if (!this._params) {
+
+    const { params } = this
+
+    if (!params) {
       throw new Error('FaceDetectionNet - load model before inference')
     }
 
@@ -42,14 +32,14 @@ export class FaceDetectionNet {
       const batchTensor = input.toBatchTensor(512, false)
 
       const x = tf.sub(tf.mul(batchTensor, tf.scalar(0.007843137718737125)), tf.scalar(1)) as tf.Tensor4D
-      const features = mobileNetV1(x, this._params.mobilenetv1_params)
+      const features = mobileNetV1(x, params.mobilenetv1)
 
       const {
         boxPredictions,
         classPredictions
-      } = predictionLayer(features.out, features.conv11, this._params.prediction_layer_params)
+      } = predictionLayer(features.out, features.conv11, params.prediction_layer)
 
-      return outputLayer(boxPredictions, classPredictions, this._params.output_layer_params)
+      return outputLayer(boxPredictions, classPredictions, params.output_layer)
     })
   }
 
@@ -91,7 +81,6 @@ export class FaceDetectionNet {
       minConfidence
     )
 
-
     const paddedHeightRelative = (netInput.getPaddings(0).y + netInput.getInputHeight(0)) / netInput.getInputHeight(0)
     const paddedWidthRelative = (netInput.getPaddings(0).x + netInput.getInputWidth(0)) / netInput.getInputWidth(0)
 
@@ -125,4 +114,12 @@ export class FaceDetectionNet {
 
     return results
   }
+
+  protected loadQuantizedParams(uri: string | undefined) {
+    return loadQuantizedParams(uri)
+  }
+
+  protected extractParams(weights: Float32Array) {
+    return extractParams(weights)
+  }
 }
\ No newline at end of file

src/faceDetectionNet/boxPredictionLayer.ts

@@ -13,11 +13,11 @@ export function boxPredictionLayer(
     const batchSize = x.shape[0]
 
     const boxPredictionEncoding = tf.reshape(
-      convLayer(x, params.box_encoding_predictor_params),
+      convLayer(x, params.box_encoding_predictor),
       [batchSize, -1, 1, 4]
     )
     const classPrediction = tf.reshape(
-      convLayer(x, params.class_predictor_params),
+      convLayer(x, params.class_predictor),
       [batchSize, -1, 3]
     )
 

src/faceDetectionNet/mobileNetV1.ts

@@ -34,13 +34,29 @@ export function mobileNetV1(x: tf.Tensor4D, params: MobileNetV1.Params) {
   return tf.tidy(() => {
 
     let conv11 = null
-    let out = pointwiseConvLayer(x, params.conv_0_params, [2, 2])
+    let out = pointwiseConvLayer(x, params.conv_0, [2, 2])
 
-    params.conv_pair_params.forEach((param, i) => {
+    const convPairParams = [
+      params.conv_1,
+      params.conv_2,
+      params.conv_3,
+      params.conv_4,
+      params.conv_5,
+      params.conv_6,
+      params.conv_7,
+      params.conv_8,
+      params.conv_9,
+      params.conv_10,
+      params.conv_11,
+      params.conv_12,
+      params.conv_13
+    ]
+
+    convPairParams.forEach((param, i) => {
       const layerIdx = i + 1
       const depthwiseConvStrides = getStridesForLayerIdx(layerIdx)
-      out = depthwiseConvLayer(out, param.depthwise_conv_params, depthwiseConvStrides)
-      out = pointwiseConvLayer(out, param.pointwise_conv_params, [1, 1])
+      out = depthwiseConvLayer(out, param.depthwise_conv, depthwiseConvStrides)
+      out = pointwiseConvLayer(out, param.pointwise_conv, [1, 1])
       if (layerIdx === 11) {
         conv11 = out
       }

src/faceDetectionNet/predictionLayer.ts

@@ -11,21 +11,21 @@ export function predictionLayer(
 ) {
   return tf.tidy(() => {
 
-    const conv0 = pointwiseConvLayer(x, params.conv_0_params, [1, 1])
-    const conv1 = pointwiseConvLayer(conv0, params.conv_1_params, [2, 2])
-    const conv2 = pointwiseConvLayer(conv1, params.conv_2_params, [1, 1])
-    const conv3 = pointwiseConvLayer(conv2, params.conv_3_params, [2, 2])
-    const conv4 = pointwiseConvLayer(conv3, params.conv_4_params, [1, 1])
-    const conv5 = pointwiseConvLayer(conv4, params.conv_5_params, [2, 2])
-    const conv6 = pointwiseConvLayer(conv5, params.conv_6_params, [1, 1])
-    const conv7 = pointwiseConvLayer(conv6, params.conv_7_params, [2, 2])
+    const conv0 = pointwiseConvLayer(x, params.conv_0, [1, 1])
+    const conv1 = pointwiseConvLayer(conv0, params.conv_1, [2, 2])
+    const conv2 = pointwiseConvLayer(conv1, params.conv_2, [1, 1])
+    const conv3 = pointwiseConvLayer(conv2, params.conv_3, [2, 2])
+    const conv4 = pointwiseConvLayer(conv3, params.conv_4, [1, 1])
+    const conv5 = pointwiseConvLayer(conv4, params.conv_5, [2, 2])
+    const conv6 = pointwiseConvLayer(conv5, params.conv_6, [1, 1])
+    const conv7 = pointwiseConvLayer(conv6, params.conv_7, [2, 2])
 
-    const boxPrediction0 = boxPredictionLayer(conv11, params.box_predictor_0_params)
-    const boxPrediction1 = boxPredictionLayer(x, params.box_predictor_1_params)
-    const boxPrediction2 = boxPredictionLayer(conv1, params.box_predictor_2_params)
-    const boxPrediction3 = boxPredictionLayer(conv3, params.box_predictor_3_params)
-    const boxPrediction4 = boxPredictionLayer(conv5, params.box_predictor_4_params)
-    const boxPrediction5 = boxPredictionLayer(conv7, params.box_predictor_5_params)
+    const boxPrediction0 = boxPredictionLayer(conv11, params.box_predictor_0)
+    const boxPrediction1 = boxPredictionLayer(x, params.box_predictor_1)
+    const boxPrediction2 = boxPredictionLayer(conv1, params.box_predictor_2)
+    const boxPrediction3 = boxPredictionLayer(conv3, params.box_predictor_3)
+    const boxPrediction4 = boxPredictionLayer(conv5, params.box_predictor_4)
+    const boxPrediction5 = boxPredictionLayer(conv7, params.box_predictor_5)
 
     const boxPredictions = tf.concat([
       boxPrediction0.boxPredictionEncoding,

src/faceDetectionNet/types.ts

@@ -18,37 +18,49 @@ export namespace MobileNetV1 {
   }
 
   export type ConvPairParams = {
-    depthwise_conv_params: DepthwiseConvParams
-    pointwise_conv_params: PointwiseConvParams
+    depthwise_conv: DepthwiseConvParams
+    pointwise_conv: PointwiseConvParams
   }
 
   export type Params = {
-    conv_0_params: PointwiseConvParams
-    conv_pair_params: ConvPairParams[]
+    conv_0: PointwiseConvParams
+    conv_1: ConvPairParams
+    conv_2: ConvPairParams
+    conv_3: ConvPairParams
+    conv_4: ConvPairParams
+    conv_5: ConvPairParams
+    conv_6: ConvPairParams
+    conv_7: ConvPairParams
+    conv_8: ConvPairParams
+    conv_9: ConvPairParams
+    conv_10: ConvPairParams
+    conv_11: ConvPairParams
+    conv_12: ConvPairParams
+    conv_13: ConvPairParams
   }
 
 }
 
 export type BoxPredictionParams = {
-  box_encoding_predictor_params: ConvParams
-  class_predictor_params: ConvParams
+  box_encoding_predictor: ConvParams
+  class_predictor: ConvParams
 }
 
 export type PredictionLayerParams = {
-  conv_0_params: PointwiseConvParams
-  conv_1_params: PointwiseConvParams
-  conv_2_params: PointwiseConvParams
-  conv_3_params: PointwiseConvParams
-  conv_4_params: PointwiseConvParams
-  conv_5_params: PointwiseConvParams
-  conv_6_params: PointwiseConvParams
-  conv_7_params: PointwiseConvParams
-  box_predictor_0_params: BoxPredictionParams
-  box_predictor_1_params: BoxPredictionParams
-  box_predictor_2_params: BoxPredictionParams
-  box_predictor_3_params: BoxPredictionParams
-  box_predictor_4_params: BoxPredictionParams
-  box_predictor_5_params: BoxPredictionParams
+  conv_0: PointwiseConvParams
+  conv_1: PointwiseConvParams
+  conv_2: PointwiseConvParams
+  conv_3: PointwiseConvParams
+  conv_4: PointwiseConvParams
+  conv_5: PointwiseConvParams
+  conv_6: PointwiseConvParams
+  conv_7: PointwiseConvParams
+  box_predictor_0: BoxPredictionParams
+  box_predictor_1: BoxPredictionParams
+  box_predictor_2: BoxPredictionParams
+  box_predictor_3: BoxPredictionParams
+  box_predictor_4: BoxPredictionParams
+  box_predictor_5: BoxPredictionParams
 }
 
 export type OutputLayerParams = {
@@ -56,7 +68,7 @@ export type OutputLayerParams = {
 }
 
 export type NetParams = {
-  mobilenetv1_params: MobileNetV1.Params,
-  prediction_layer_params: PredictionLayerParams,
-  output_layer_params: OutputLayerParams
+  mobilenetv1: MobileNetV1.Params,
+  prediction_layer: PredictionLayerParams,
+  output_layer: OutputLayerParams
 }

src/faceLandmarkNet/FaceLandmarkNet.ts

@@ -24,36 +24,13 @@ function maxPool(x: tf.Tensor4D, strides: [number, number] = [2, 2]): tf.Tensor4
 
 export class FaceLandmarkNet extends NeuralNetwork<NetParams> {
 
-  public async load(weightsOrUrl: Float32Array | string | undefined): Promise<void> {
-    if (weightsOrUrl instanceof Float32Array) {
-      this.extractWeights(weightsOrUrl)
-      return
-    }
-
-    if (weightsOrUrl && typeof weightsOrUrl !== 'string') {
-      throw new Error('FaceLandmarkNet.load - expected model uri, or weights as Float32Array')
-    }
-    const {
-      paramMappings,
-      params
-    } = await loadQuantizedParams(weightsOrUrl)
-
-    this._paramMappings = paramMappings
-    this._params = params
-  }
-
-  public extractWeights(weights: Float32Array) {
-    const {
-      paramMappings,
-      params
-    } = extractParams(weights)
-
-    this._paramMappings = paramMappings
-    this._params = params
+  constructor() {
+    super('FaceLandmarkNet')
   }
 
   public forwardInput(input: NetInput): tf.Tensor2D {
-    const params = this._params
+
+    const { params } = this
 
     if (!params) {
       throw new Error('FaceLandmarkNet - load model before inference')
@@ -62,20 +39,20 @@ export class FaceLandmarkNet extends NeuralNetwork<NetParams> {
     return tf.tidy(() => {
       const batchTensor = input.toBatchTensor(128, true)
 
-      let out = conv(batchTensor, params.conv0_params)
+      let out = conv(batchTensor, params.conv0)
       out = maxPool(out)
-      out = conv(out, params.conv1_params)
-      out = conv(out, params.conv2_params)
+      out = conv(out, params.conv1)
+      out = conv(out, params.conv2)
       out = maxPool(out)
-      out = conv(out, params.conv3_params)
-      out = conv(out, params.conv4_params)
+      out = conv(out, params.conv3)
+      out = conv(out, params.conv4)
       out = maxPool(out)
-      out = conv(out, params.conv5_params)
-      out = conv(out, params.conv6_params)
+      out = conv(out, params.conv5)
+      out = conv(out, params.conv6)
       out = maxPool(out, [1, 1])
-      out = conv(out, params.conv7_params)
-      const fc0 = tf.relu(fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc0_params))
-      const fc1 = fullyConnectedLayer(fc0, params.fc1_params)
+      out = conv(out, params.conv7)
+      const fc0 = tf.relu(fullyConnectedLayer(out.as2D(out.shape[0], -1), params.fc0))
+      const fc1 = fullyConnectedLayer(fc0, params.fc1)
 
       const createInterleavedTensor = (fillX: number, fillY: number) =>
         tf.stack([
@@ -145,4 +122,12 @@ export class FaceLandmarkNet extends NeuralNetwork<NetParams> {
       ? landmarksForBatch
       : landmarksForBatch[0]
   }
+
+  protected loadQuantizedParams(uri: string | undefined) {
+    return loadQuantizedParams(uri)
+  }
+
+  protected extractParams(weights: Float32Array) {
+    return extractParams(weights)
+  }
 }
\ No newline at end of file

src/faceLandmarkNet/extractParams.ts

 import * as tf from '@tensorflow/tfjs-core';
 
-import { extractConvParamsFactory } from '../commons/extractConvParamsFactory';
 import { extractWeightsFactory } from '../commons/extractWeightsFactory';
-import { ParamMapping } from '../commons/types';
+import { ConvParams, ParamMapping } from '../commons/types';
 import { FCParams, NetParams } from './types';
 
 export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
+
   const paramMappings: ParamMapping[] = []
 
   const {
@@ -13,9 +13,29 @@ export function extractParams(weights: Float32Array): { params: NetParams, param
     getRemainingWeights
   } = extractWeightsFactory(weights)
 
-  const extractConvParams = extractConvParamsFactory(extractWeights, paramMappings)
+  function extractConvParams(
+    channelsIn: number,
+    channelsOut: number,
+    filterSize: number,
+    mappedPrefix: string
+  ): ConvParams {
+
+    const filters = tf.tensor4d(
+      extractWeights(channelsIn * channelsOut * filterSize * filterSize),
+      [filterSize, filterSize, channelsIn, channelsOut]
+    )
+    const bias = tf.tensor1d(extractWeights(channelsOut))
+
+    paramMappings.push(
+      { paramPath: `${mappedPrefix}/filters` },
+      { paramPath: `${mappedPrefix}/bias` }
+    )
+
+    return { filters, bias }
+  }
 
   function extractFcParams(channelsIn: number, channelsOut: number, mappedPrefix: string): FCParams {
+
     const fc_weights = tf.tensor2d(extractWeights(channelsIn * channelsOut), [channelsIn, channelsOut])
     const fc_bias = tf.tensor1d(extractWeights(channelsOut))
 
@@ -30,16 +50,16 @@ export function extractParams(weights: Float32Array): { params: NetParams, param
     }
   }
 
-  const conv0_params = extractConvParams(3, 32, 3, 'conv0_params')
-  const conv1_params = extractConvParams(32, 64, 3, 'conv1_params')
-  const conv2_params = extractConvParams(64, 64, 3, 'conv2_params')
-  const conv3_params = extractConvParams(64, 64, 3, 'conv3_params')
-  const conv4_params = extractConvParams(64, 64, 3, 'conv4_params')
-  const conv5_params = extractConvParams(64, 128, 3, 'conv5_params')
-  const conv6_params = extractConvParams(128, 128, 3, 'conv6_params')
-  const conv7_params = extractConvParams(128, 256, 3, 'conv7_params')
-  const fc0_params = extractFcParams(6400, 1024, 'fc0_params')
-  const fc1_params = extractFcParams(1024, 136, 'fc1_params')
+  const conv0 = extractConvParams(3, 32, 3, 'conv0')
+  const conv1 = extractConvParams(32, 64, 3, 'conv1')
+  const conv2 = extractConvParams(64, 64, 3, 'conv2')
+  const conv3 = extractConvParams(64, 64, 3, 'conv3')
+  const conv4 = extractConvParams(64, 64, 3, 'conv4')
+  const conv5 = extractConvParams(64, 128, 3, 'conv5')
+  const conv6 = extractConvParams(128, 128, 3, 'conv6')
+  const conv7 = extractConvParams(128, 256, 3, 'conv7')
+  const fc0 = extractFcParams(6400, 1024, 'fc0')
+  const fc1 = extractFcParams(1024, 136, 'fc1')
 
   if (getRemainingWeights().length !== 0) {
     throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
@@ -48,16 +68,16 @@ export function extractParams(weights: Float32Array): { params: NetParams, param
   return {
     paramMappings,
     params: {
-      conv0_params,
-      conv1_params,
-      conv2_params,
-      conv3_params,
-      conv4_params,
-      conv5_params,
-      conv6_params,
-      conv7_params,
-      fc0_params,
-      fc1_params
+      conv0,
+      conv1,
+      conv2,
+      conv3,
+      conv4,
+      conv5,
+      conv6,
+      conv7,
+      fc0,
+      fc1
     }
   }
 }
\ No newline at end of file

src/faceLandmarkNet/loadQuantizedParams.ts

 import * as tf from '@tensorflow/tfjs-core';
 
-import { extractWeightEntry } from '../commons/extractWeightEntry';
+import { disposeUnusedWeightTensors } from '../commons/disposeUnusedWeightTensors';
+import { extractWeightEntryFactory } from '../commons/extractWeightEntryFactory';
 import { loadWeightMap } from '../commons/loadWeightMap';
 import { ConvParams, ParamMapping } from '../commons/types';
 import { FCParams, NetParams } from './types';
@@ -9,30 +10,20 @@ const DEFAULT_MODEL_NAME = 'face_landmark_68_model'
 
 function extractorsFactory(weightMap: any, paramMappings: ParamMapping[]) {
 
+  const extractWeightEntry = extractWeightEntryFactory(weightMap, paramMappings)
+
   function extractConvParams(prefix: string, mappedPrefix: string): ConvParams {
-    const filtersEntry = extractWeightEntry(weightMap, `${prefix}/kernel`, 4)
-    const biasEntry = extractWeightEntry(weightMap, `${prefix}/bias`, 1)
-    paramMappings.push(
-      { originalPath: filtersEntry.path, paramPath: `${mappedPrefix}/filters` },
-      { originalPath: biasEntry.path, paramPath: `${mappedPrefix}/bias` }
-    )
-    return {
-      filters: filtersEntry.tensor as tf.Tensor4D,
-      bias: biasEntry.tensor as tf.Tensor1D
-    }
+    const filters = extractWeightEntry<tf.Tensor4D>(`${prefix}/kernel`, 4, `${mappedPrefix}/filters`)
+    const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1, `${mappedPrefix}/bias`)
+
+    return { filters, bias }
   }
 
   function extractFcParams(prefix: string, mappedPrefix: string): FCParams {
-    const weightsEntry = extractWeightEntry(weightMap, `${prefix}/kernel`, 2)
-    const biasEntry = extractWeightEntry(weightMap, `${prefix}/bias`, 1)
-    paramMappings.push(
-      { originalPath: weightsEntry.path, paramPath: `${mappedPrefix}/weights` },
-      { originalPath: biasEntry.path, paramPath: `${mappedPrefix}/bias` }
-    )
-    return {
-      weights: weightsEntry.tensor as tf.Tensor2D,
-      bias: biasEntry.tensor as tf.Tensor1D
-    }
+    const weights = extractWeightEntry<tf.Tensor2D>(`${prefix}/kernel`, 2, `${mappedPrefix}/weights`)
+    const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1, `${mappedPrefix}/bias`)
+
+    return { weights, bias }
   }
 
   return {
@@ -54,17 +45,19 @@ export async function loadQuantizedParams(
   } = extractorsFactory(weightMap, paramMappings)
 
   const params = {
-    conv0_params: extractConvParams('conv2d_0', 'conv0_params'),
-    conv1_params: extractConvParams('conv2d_1', 'conv1_params'),
-    conv2_params: extractConvParams('conv2d_2', 'conv2_params'),
-    conv3_params: extractConvParams('conv2d_3', 'conv3_params'),
-    conv4_params: extractConvParams('conv2d_4', 'conv4_params'),
-    conv5_params: extractConvParams('conv2d_5', 'conv5_params'),
-    conv6_params: extractConvParams('conv2d_6', 'conv6_params'),
-    conv7_params: extractConvParams('conv2d_7', 'conv7_params'),
-    fc0_params: extractFcParams('dense', 'fc0_params'),
-    fc1_params: extractFcParams('logits', 'fc1_params')
+    conv0: extractConvParams('conv2d_0', 'conv0'),
+    conv1: extractConvParams('conv2d_1', 'conv1'),
+    conv2: extractConvParams('conv2d_2', 'conv2'),
+    conv3: extractConvParams('conv2d_3', 'conv3'),
+    conv4: extractConvParams('conv2d_4', 'conv4'),
+    conv5: extractConvParams('conv2d_5', 'conv5'),
+    conv6: extractConvParams('conv2d_6', 'conv6'),
+    conv7: extractConvParams('conv2d_7', 'conv7'),
+    fc0: extractFcParams('dense', 'fc0'),
+    fc1: extractFcParams('logits', 'fc1')
   }
 
+  disposeUnusedWeightTensors(weightMap, paramMappings)
+
   return { params, paramMappings }
 }
\ No newline at end of file

src/faceLandmarkNet/types.ts

@@ -8,14 +8,14 @@ export type FCParams = {
 }
 
 export type NetParams = {
-  conv0_params: ConvParams
-  conv1_params: ConvParams
-  conv2_params: ConvParams
-  conv3_params: ConvParams
-  conv4_params: ConvParams
-  conv5_params: ConvParams
-  conv6_params: ConvParams
-  conv7_params: ConvParams
-  fc0_params: FCParams
-  fc1_params: FCParams
+  conv0: ConvParams
+  conv1: ConvParams
+  conv2: ConvParams
+  conv3: ConvParams
+  conv4: ConvParams
+  conv5: ConvParams
+  conv6: ConvParams
+  conv7: ConvParams
+  fc0: FCParams
+  fc1: FCParams
 }
\ No newline at end of file

src/faceRecognitionNet/FaceRecognitionNet.ts

 import * as tf from '@tensorflow/tfjs-core';
 
+import { NeuralNetwork } from '../commons/NeuralNetwork';
 import { NetInput } from '../NetInput';
 import { toNetInput } from '../toNetInput';
 import { TNetInput } from '../types';
@@ -10,28 +11,17 @@ import { normalize } from './normalize';
 import { residual, residualDown } from './residualLayer';
 import { NetParams } from './types';
 
-export class FaceRecognitionNet {
+export class FaceRecognitionNet extends NeuralNetwork<NetParams> {
 
-  private _params: NetParams
-
-  public async load(weightsOrUrl: Float32Array | string | undefined): Promise<void> {
-    if (weightsOrUrl instanceof Float32Array) {
-      this.extractWeights(weightsOrUrl)
-      return
-    }
-
-    if (weightsOrUrl && typeof weightsOrUrl !== 'string') {
-      throw new Error('FaceLandmarkNet.load - expected model uri, or weights as Float32Array')
-    }
-    this._params = await loadQuantizedParams(weightsOrUrl)
-  }
-
-  public extractWeights(weights: Float32Array) {
-    this._params = extractParams(weights)
+  constructor() {
+    super('FaceRecognitionNet')
   }
 
   public forwardInput(input: NetInput): tf.Tensor2D {
-    if (!this._params) {
+
+    const { params } = this
+
+    if (!params) {
       throw new Error('FaceRecognitionNet - load model before inference')
     }
 
@@ -40,29 +30,29 @@ export class FaceRecognitionNet {
 
       const normalized = normalize(batchTensor)
 
-      let out = convDown(normalized, this._params.conv32_down)
+      let out = convDown(normalized, params.conv32_down)
       out = tf.maxPool(out, 3, 2, 'valid')
 
-      out = residual(out, this._params.conv32_1)
-      out = residual(out, this._params.conv32_2)
-      out = residual(out, this._params.conv32_3)
+      out = residual(out, params.conv32_1)
+      out = residual(out, params.conv32_2)
+      out = residual(out, params.conv32_3)
 
-      out = residualDown(out, this._params.conv64_down)
-      out = residual(out, this._params.conv64_1)
-      out = residual(out, this._params.conv64_2)
-      out = residual(out, this._params.conv64_3)
+      out = residualDown(out, params.conv64_down)
+      out = residual(out, params.conv64_1)
+      out = residual(out, params.conv64_2)
+      out = residual(out, params.conv64_3)
 
-      out = residualDown(out, this._params.conv128_down)
-      out = residual(out, this._params.conv128_1)
-      out = residual(out, this._params.conv128_2)
+      out = residualDown(out, params.conv128_down)
+      out = residual(out, params.conv128_1)
+      out = residual(out, params.conv128_2)
 
-      out = residualDown(out, this._params.conv256_down)
-      out = residual(out, this._params.conv256_1)
-      out = residual(out, this._params.conv256_2)
-      out = residualDown(out, this._params.conv256_down_out)
+      out = residualDown(out, params.conv256_down)
+      out = residual(out, params.conv256_1)
+      out = residual(out, params.conv256_2)
+      out = residualDown(out, params.conv256_down_out)
 
       const globalAvg = out.mean([1, 2]) as tf.Tensor2D
-      const fullyConnected = tf.matMul(globalAvg, this._params.fc)
+      const fullyConnected = tf.matMul(globalAvg, params.fc)
 
       return fullyConnected
     })
@@ -89,4 +79,12 @@ export class FaceRecognitionNet {
       ? faceDescriptorsForBatch
       : faceDescriptorsForBatch[0]
   }
+
+  protected loadQuantizedParams(uri: string | undefined) {
+    return loadQuantizedParams(uri)
+  }
+
+  protected extractParams(weights: Float32Array) {
+    return extractParams(weights)
+  }
 }
\ No newline at end of file

src/faceRecognitionNet/extractParams.ts

 import * as tf from '@tensorflow/tfjs-core';
 
 import { extractWeightsFactory } from '../commons/extractWeightsFactory';
-import { ExtractWeightsFunction } from '../commons/types';
+import { ConvParams, ExtractWeightsFunction, ParamMapping } from '../commons/types';
 import { isFloat } from '../utils';
 import { ConvLayerParams, NetParams, ResidualLayerParams, ScaleLayerParams } from './types';
 
-function extractorsFactory(extractWeights: ExtractWeightsFunction) {
+function extractorsFactory(extractWeights: ExtractWeightsFunction, paramMappings: ParamMapping[]) {
 
   function extractFilterValues(numFilterValues: number, numFilters: number, filterSize: number): tf.Tensor4D {
     const weights = extractWeights(numFilterValues)
@@ -15,15 +15,42 @@ function extractorsFactory(extractWeights: ExtractWeightsFunction) {
       throw new Error(`depth has to be an integer: ${depth}, weights.length: ${weights.length}, numFilters: ${numFilters}, filterSize: ${filterSize}`)
     }
 
-    return tf.transpose(
-      tf.tensor4d(weights, [numFilters, depth, filterSize, filterSize]),
-      [2, 3, 1, 0]
+    return tf.tidy(
+      () => tf.transpose(
+        tf.tensor4d(weights, [numFilters, depth, filterSize, filterSize]),
+        [2, 3, 1, 0]
+      )
     )
   }
 
-  function extractScaleLayerParams(numWeights: number): ScaleLayerParams {
+  function extractConvParams(
+    numFilterValues: number,
+    numFilters: number,
+    filterSize: number,
+    mappedPrefix: string
+  ): ConvParams {
+
+    const filters = extractFilterValues(numFilterValues, numFilters, filterSize)
+    const bias = tf.tensor1d(extractWeights(numFilters))
+
+    paramMappings.push(
+      { paramPath: `${mappedPrefix}/filters` },
+      { paramPath: `${mappedPrefix}/bias` }
+    )
+
+    return { filters, bias }
+  }
+
+  function extractScaleLayerParams(numWeights: number, mappedPrefix: string): ScaleLayerParams {
+
     const weights = tf.tensor1d(extractWeights(numWeights))
     const biases = tf.tensor1d(extractWeights(numWeights))
+
+    paramMappings.push(
+      { paramPath: `${mappedPrefix}/weights` },
+      { paramPath: `${mappedPrefix}/biases` }
+    )
+
     return {
       weights,
       biases
@@ -33,34 +60,28 @@ function extractorsFactory(extractWeights: ExtractWeightsFunction) {
   function extractConvLayerParams(
     numFilterValues: number,
     numFilters: number,
-    filterSize: number
+    filterSize: number,
+    mappedPrefix: string
   ): ConvLayerParams {
-    const conv_filters = extractFilterValues(numFilterValues, numFilters, filterSize)
-    const conv_bias = tf.tensor1d(extractWeights(numFilters))
-    const scale = extractScaleLayerParams(numFilters)
 
-    return {
-      conv: {
-        filters: conv_filters,
-        bias: conv_bias
-      },
-      scale
-    }
+    const conv = extractConvParams(numFilterValues, numFilters, filterSize, `${mappedPrefix}/conv`)
+    const scale = extractScaleLayerParams(numFilters, `${mappedPrefix}/scale`)
+
+    return { conv, scale }
   }
 
   function extractResidualLayerParams(
     numFilterValues: number,
     numFilters: number,
     filterSize: number,
+    mappedPrefix: string,
     isDown: boolean = false
   ): ResidualLayerParams {
-    const conv1: ConvLayerParams = extractConvLayerParams((isDown ? 0.5 : 1) * numFilterValues, numFilters, filterSize)
-    const conv2: ConvLayerParams = extractConvLayerParams(numFilterValues, numFilters, filterSize)
 
-    return {
-      conv1,
-      conv2
-    }
+    const conv1 = extractConvLayerParams((isDown ? 0.5 : 1) * numFilterValues, numFilters, filterSize, `${mappedPrefix}/conv1`)
+    const conv2 = extractConvLayerParams(numFilterValues, numFilters, filterSize, `${mappedPrefix}/conv2`)
+
+    return { conv1, conv2 }
   }
 
   return {
@@ -70,43 +91,49 @@ function extractorsFactory(extractWeights: ExtractWeightsFunction) {
 
 }
 
-export function extractParams(weights: Float32Array): NetParams {
+export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
+
   const {
     extractWeights,
     getRemainingWeights
   } = extractWeightsFactory(weights)
 
+  const paramMappings: ParamMapping[] = []
+
   const {
     extractConvLayerParams,
     extractResidualLayerParams
-  } = extractorsFactory(extractWeights)
+  } = extractorsFactory(extractWeights, paramMappings)
 
-  const conv32_down = extractConvLayerParams(4704, 32, 7)
-  const conv32_1 = extractResidualLayerParams(9216, 32, 3)
-  const conv32_2 = extractResidualLayerParams(9216, 32, 3)
-  const conv32_3 = extractResidualLayerParams(9216, 32, 3)
+  const conv32_down = extractConvLayerParams(4704, 32, 7, 'conv32_down')
+  const conv32_1 = extractResidualLayerParams(9216, 32, 3, 'conv32_1')
+  const conv32_2 = extractResidualLayerParams(9216, 32, 3, 'conv32_2')
+  const conv32_3 = extractResidualLayerParams(9216, 32, 3, 'conv32_3')
 
-  const conv64_down = extractResidualLayerParams(36864, 64, 3, true)
-  const conv64_1 = extractResidualLayerParams(36864, 64, 3)
-  const conv64_2 = extractResidualLayerParams(36864, 64, 3)
-  const conv64_3 = extractResidualLayerParams(36864, 64, 3)
+  const conv64_down = extractResidualLayerParams(36864, 64, 3, 'conv64_down', true)
+  const conv64_1 = extractResidualLayerParams(36864, 64, 3, 'conv64_1')
+  const conv64_2 = extractResidualLayerParams(36864, 64, 3, 'conv64_2')
+  const conv64_3 = extractResidualLayerParams(36864, 64, 3, 'conv64_3')
 
-  const conv128_down = extractResidualLayerParams(147456, 128, 3, true)
-  const conv128_1 = extractResidualLayerParams(147456, 128, 3)
-  const conv128_2 = extractResidualLayerParams(147456, 128, 3)
+  const conv128_down = extractResidualLayerParams(147456, 128, 3, 'conv128_down', true)
+  const conv128_1 = extractResidualLayerParams(147456, 128, 3, 'conv128_1')
+  const conv128_2 = extractResidualLayerParams(147456, 128, 3, 'conv128_2')
 
-  const conv256_down = extractResidualLayerParams(589824, 256, 3, true)
-  const conv256_1 = extractResidualLayerParams(589824, 256, 3)
-  const conv256_2 = extractResidualLayerParams(589824, 256, 3)
-  const conv256_down_out = extractResidualLayerParams(589824, 256, 3)
+  const conv256_down = extractResidualLayerParams(589824, 256, 3, 'conv256_down', true)
+  const conv256_1 = extractResidualLayerParams(589824, 256, 3, 'conv256_1')
+  const conv256_2 = extractResidualLayerParams(589824, 256, 3, 'conv256_2')
+  const conv256_down_out = extractResidualLayerParams(589824, 256, 3, 'conv256_down_out')
 
-  const fc = tf.transpose(tf.tensor2d(extractWeights(256 * 128), [128, 256]), [1, 0])
+  const fc = tf.tidy(
+    () => tf.transpose(tf.tensor2d(extractWeights(256 * 128), [128, 256]), [1, 0])
+  )
+  paramMappings.push({ paramPath: `fc` })
 
   if (getRemainingWeights().length !== 0) {
     throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
   }
 
-  return {
+  const params = {
     conv32_down,
     conv32_1,
     conv32_2,
@@ -124,4 +151,6 @@ export function extractParams(weights: Float32Array): NetParams {
     conv256_down_out,
     fc
   }
+
+  return { params, paramMappings }
 }
\ No newline at end of file

src/faceRecognitionNet/loadQuantizedParams.ts

-import { isTensor1D, isTensor2D, isTensor4D } from '../commons/isTensor';
+import * as tf from '@tensorflow/tfjs-core';
+
+import { disposeUnusedWeightTensors } from '../commons/disposeUnusedWeightTensors';
+import { extractWeightEntryFactory } from '../commons/extractWeightEntryFactory';
+import { isTensor2D } from '../commons/isTensor';
 import { loadWeightMap } from '../commons/loadWeightMap';
-import { ConvLayerParams, ResidualLayerParams, ScaleLayerParams } from './types';
+import { ParamMapping } from '../commons/types';
+import { ConvLayerParams, NetParams, ResidualLayerParams, ScaleLayerParams } from './types';
 
 const DEFAULT_MODEL_NAME = 'face_recognition_model'
 
-function extractorsFactory(weightMap: any) {
+function extractorsFactory(weightMap: any, paramMappings: ParamMapping[]) {
 
-  function extractScaleLayerParams(prefix: string): ScaleLayerParams {
-    const params = {
-      weights: weightMap[`${prefix}/scale/weights`],
-      biases: weightMap[`${prefix}/scale/biases`]
-    }
+  const extractWeightEntry = extractWeightEntryFactory(weightMap, paramMappings)
 
-    if (!isTensor1D(params.weights)) {
-      throw new Error(`expected weightMap[${prefix}/scale/weights] to be a Tensor1D, instead have ${params.weights}`)
-    }
+  function extractScaleLayerParams(prefix: string): ScaleLayerParams {
 
-    if (!isTensor1D(params.biases)) {
-      throw new Error(`expected weightMap[${prefix}/scale/biases] to be a Tensor1D, instead have ${params.biases}`)
-    }
+    const weights = extractWeightEntry<tf.Tensor1D>(`${prefix}/scale/weights`, 1)
+    const biases = extractWeightEntry<tf.Tensor1D>(`${prefix}/scale/biases`, 1)
 
-    return params
+    return { weights, biases }
   }
 
   function extractConvLayerParams(prefix: string): ConvLayerParams {
-    const params = {
-      filters: weightMap[`${prefix}/conv/filters`],
-      bias: weightMap[`${prefix}/conv/bias`]
-    }
-
-    if (!isTensor4D(params.filters)) {
-      throw new Error(`expected weightMap[${prefix}/conv/filters] to be a Tensor1D, instead have ${params.filters}`)
-    }
 
-    if (!isTensor1D(params.bias)) {
-      throw new Error(`expected weightMap[${prefix}/conv/bias] to be a Tensor1D, instead have ${params.bias}`)
-    }
+    const filters = extractWeightEntry<tf.Tensor4D>(`${prefix}/conv/filters`, 4)
+    const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/conv/bias`, 1)
+    const scale = extractScaleLayerParams(prefix)
 
-    return {
-      conv: params,
-      scale: extractScaleLayerParams(prefix)
-    }
+    return { conv: { filters, bias }, scale }
   }
 
   function extractResidualLayerParams(prefix: string): ResidualLayerParams {
@@ -57,13 +44,17 @@ function extractorsFactory(weightMap: any) {
 
 }
 
-export async function loadQuantizedParams(uri: string | undefined): Promise<any> {
+export async function loadQuantizedParams(
+  uri: string | undefined
+): Promise<{ params: NetParams, paramMappings: ParamMapping[] }> {
+
   const weightMap = await loadWeightMap(uri, DEFAULT_MODEL_NAME)
+  const paramMappings: ParamMapping[] = []
 
   const {
     extractConvLayerParams,
     extractResidualLayerParams
-  } = extractorsFactory(weightMap)
+  } = extractorsFactory(weightMap, paramMappings)
 
   const conv32_down = extractConvLayerParams('conv32_down')
   const conv32_1 = extractResidualLayerParams('conv32_1')
@@ -85,12 +76,13 @@ export async function loadQuantizedParams(uri: string | undefined): Promise<any>
   const conv256_down_out = extractResidualLayerParams('conv256_down_out')
 
   const fc = weightMap['fc']
+  paramMappings.push({ originalPath: 'fc', paramPath: 'fc' })
 
   if (!isTensor2D(fc)) {
     throw new Error(`expected weightMap[fc] to be a Tensor2D, instead have ${fc}`)
   }
 
-  return {
+  const params = {
     conv32_down,
     conv32_1,
     conv32_2,
@@ -108,4 +100,8 @@ export async function loadQuantizedParams(uri: string | undefined): Promise<any>
     conv256_down_out,
     fc
   }
+
+  disposeUnusedWeightTensors(weightMap, paramMappings)
+
+  return { params, paramMappings }
 }
\ No newline at end of file

test/tests/commons/NeuralNetwork.test.ts

-import { NeuralNetwork } from '../../../src/commons/NeuralNetwork';
 import * as tf from '@tensorflow/tfjs-core';
 
+import { NeuralNetwork } from '../../../src/commons/NeuralNetwork';
+
 class FakeNeuralNetwork extends NeuralNetwork<any> {
   constructor(
     convFilter: tf.Tensor = tf.tensor(0),
     convBias: tf.Tensor = tf.tensor(0),
     fcWeights: tf.Tensor = tf.tensor(0)
   ) {
-    super()
+    super('FakeNeuralNetwork')
     this._params = {
       conv: {
         filter: convFilter,

test/tests/e2e/faceDetectionNet.test.ts

 import * as faceapi from '../../../src';
 import { FaceDetection } from '../../../src/faceDetectionNet/FaceDetection';
 import { IRect } from '../../../src/Rect';
-import { expectMaxDelta } from '../../utils';
+import { expectAllTensorsReleased, expectMaxDelta } from '../../utils';
 
 function expectRectClose(
   result: IRect,
@@ -110,4 +110,33 @@ describe('faceDetectionNet', () => {
 
   })
 
+  describe('no memory leaks', () => {
+
+    describe('NeuralNetwork, uncompressed model', () => {
+
+      it('disposes all param tensors', async () => {
+        await expectAllTensorsReleased(async () => {
+          const res = await fetch('base/weights/uncompressed/face_detection_model.weights')
+          const weights = new Float32Array(await res.arrayBuffer())
+          const net = faceapi.faceDetectionNet(weights)
+          net.dispose()
+        })
+      })
+
+    })
+
+    describe('NeuralNetwork, quantized model', () => {
+
+      it('disposes all param tensors', async () => {
+        await expectAllTensorsReleased(async () => {
+          const net = new faceapi.FaceDetectionNet()
+          await net.load('base/weights')
+          net.dispose()
+        })
+      })
+
+    })
+
+  })
+
 })
\ No newline at end of file

test/tests/e2e/faceLandmarkNet.test.ts

@@ -238,6 +238,31 @@ describe('faceLandmarkNet', () => {
       await faceLandmarkNet.load('base/weights')
     })
 
+    describe('NeuralNetwork, uncompressed model', () => {
+
+      it('disposes all param tensors', async () => {
+        await expectAllTensorsReleased(async () => {
+          const res = await fetch('base/weights/uncompressed/face_landmark_68_model.weights')
+          const weights = new Float32Array(await res.arrayBuffer())
+          const net = faceapi.faceLandmarkNet(weights)
+          net.dispose()
+        })
+      })
+
+    })
+
+    describe('NeuralNetwork, quantized model', () => {
+
+      it('disposes all param tensors', async () => {
+        await expectAllTensorsReleased(async () => {
+          const net = new faceapi.FaceLandmarkNet()
+          await net.load('base/weights')
+          net.dispose()
+        })
+      })
+
+    })
+
     describe('forwardInput', () => {
 
       it('single image element', async () => {

test/tests/e2e/faceRecognitionNet.test.ts

@@ -166,6 +166,35 @@ describe('faceRecognitionNet', () => {
       faceRecognitionNet = faceapi.faceRecognitionNet(weights)
     })
 
+    afterAll(async () => {
+      faceRecognitionNet.dispose()
+    })
+
+    describe('NeuralNetwork, uncompressed model', () => {
+
+      it('disposes all param tensors', async () => {
+        await expectAllTensorsReleased(async () => {
+          const res = await fetch('base/weights/uncompressed/face_recognition_model.weights')
+          const weights = new Float32Array(await res.arrayBuffer())
+          const net = faceapi.faceRecognitionNet(weights)
+          net.dispose()
+        })
+      })
+
+    })
+
+    describe('NeuralNetwork, quantized model', () => {
+
+      it('disposes all param tensors', async () => {
+        await expectAllTensorsReleased(async () => {
+          const net = new faceapi.FaceRecognitionNet()
+          await net.load('base/weights')
+          net.dispose()
+        })
+      })
+
+    })
+
     describe('forwardInput', () => {
 
       it('single image element', async () => {
@@ -292,5 +321,4 @@ describe('faceRecognitionNet', () => {
     })
   })
 
-
 })
\ No newline at end of file