Skip to content

F
face
Commit 9fef1428 by vincent
Options

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

parent a6a68a56
Showing with 718 additions and 519 deletions
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 0 → 100644
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 deleted 100644 → 0
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 deleted 100644 → 0
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 0 → 100644
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
constructor() {
super('FaceDetectionNet')
}
if (weightsOrUrl && typeof weightsOrUrl !== 'string') {
throw new Error('FaceDetectionNet.load - expected model uri, or weights as Float32Array')
}
this._params = await loadQuantizedParams(weightsOrUrl)
}
public forwardInput(input: NetInput) {
public extractWeights(weights: Float32Array) {
this._params = extractParams(weights)
}
const { params } = this
public forwardInput(input: NetInput) {
if (!this._params) {
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/extractParams.ts
import * as tf from '@tensorflow/tfjs-core';
import { extractWeightsFactory } from '../commons/extractWeightsFactory';
import { ConvParams } from '../commons/types';
import { ConvParams, ExtractWeightsFunction, ParamMapping } from '../commons/types';
import { MobileNetV1, NetParams, PointwiseConvParams, PredictionLayerParams } from './types';
function extractorsFactory(extractWeights: (numWeights: number) => Float32Array) {
function extractorsFactory(extractWeights: ExtractWeightsFunction, paramMappings: ParamMapping[]) {
function extractDepthwiseConvParams(numChannels: number, mappedPrefix: string): MobileNetV1.DepthwiseConvParams {
function extractDepthwiseConvParams(numChannels: number): MobileNetV1.DepthwiseConvParams {
const filters = tf.tensor4d(extractWeights(3 * 3 * numChannels), [3, 3, numChannels, 1])
const batch_norm_scale = tf.tensor1d(extractWeights(numChannels))
const batch_norm_offset = tf.tensor1d(extractWeights(numChannels))
const batch_norm_mean = tf.tensor1d(extractWeights(numChannels))
const batch_norm_variance = tf.tensor1d(extractWeights(numChannels))
paramMappings.push(
{ paramPath: `${mappedPrefix}/filters` },
{ paramPath: `${mappedPrefix}/batch_norm_scale` },
{ paramPath: `${mappedPrefix}/batch_norm_offset` },
{ paramPath: `${mappedPrefix}/batch_norm_mean` },
{ paramPath: `${mappedPrefix}/batch_norm_variance` }
)
return {
filters,
batch_norm_scale,
......@@ -25,29 +34,36 @@ function extractorsFactory(extractWeights: (numWeights: number) => Float32Array)
function extractConvParams(
channelsIn: number,
channelsOut: number,
filterSize: number
filterSize: number,
mappedPrefix: string,
isPointwiseConv?: boolean
): ConvParams {
const filters = tf.tensor4d(
extractWeights(channelsIn * channelsOut * filterSize * filterSize),
[filterSize, filterSize, channelsIn, channelsOut]
)
const bias = tf.tensor1d(extractWeights(channelsOut))
return {
filters,
bias
}
paramMappings.push(
{ paramPath: `${mappedPrefix}/filters` },
{ paramPath: `${mappedPrefix}/${isPointwiseConv ? 'batch_norm_offset' : 'bias'}` }
)
return { filters, bias }
}
function extractPointwiseConvParams(
channelsIn: number,
channelsOut: number,
filterSize: number
filterSize: number,
mappedPrefix: string
): PointwiseConvParams {
const {
filters,
bias
} = extractConvParams(channelsIn, channelsOut, filterSize)
} = extractConvParams(channelsIn, channelsOut, filterSize, mappedPrefix, true)
return {
filters,
......@@ -57,115 +73,118 @@ function extractorsFactory(extractWeights: (numWeights: number) => Float32Array)
function extractConvPairParams(
channelsIn: number,
channelsOut: number
channelsOut: number,
mappedPrefix: string
): MobileNetV1.ConvPairParams {
const depthwise_conv_params = extractDepthwiseConvParams(channelsIn)
const pointwise_conv_params = extractPointwiseConvParams(channelsIn, channelsOut, 1)
return {
depthwise_conv_params,
pointwise_conv_params
}
const depthwise_conv = extractDepthwiseConvParams(channelsIn, `${mappedPrefix}/depthwise_conv`)
const pointwise_conv = extractPointwiseConvParams(channelsIn, channelsOut, 1, `${mappedPrefix}/pointwise_conv`)
return { depthwise_conv, pointwise_conv }
}
function extractMobilenetV1Params(): MobileNetV1.Params {
const conv_0_params = extractPointwiseConvParams(3, 32, 3)
const channelNumPairs = [
[32, 64],
[64, 128],
[128, 128],
[128, 256],
[256, 256],
[256, 512],
[512, 512],
[512, 512],
[512, 512],
[512, 512],
[512, 512],
[512, 1024],
[1024, 1024]
]
const conv_pair_params = channelNumPairs.map(
([channelsIn, channelsOut]) => extractConvPairParams(channelsIn, channelsOut)
)
const conv_0 = extractPointwiseConvParams(3, 32, 3, 'mobilenetv1/conv_0')
const conv_1 = extractConvPairParams(32, 64, 'mobilenetv1/conv_1')
const conv_2 = extractConvPairParams(64, 128, 'mobilenetv1/conv_2')
const conv_3 = extractConvPairParams(128, 128, 'mobilenetv1/conv_3')
const conv_4 = extractConvPairParams(128, 256, 'mobilenetv1/conv_4')
const conv_5 = extractConvPairParams(256, 256, 'mobilenetv1/conv_5')
const conv_6 = extractConvPairParams(256, 512, 'mobilenetv1/conv_6')
const conv_7 = extractConvPairParams(512, 512, 'mobilenetv1/conv_7')
const conv_8 = extractConvPairParams(512, 512, 'mobilenetv1/conv_8')
const conv_9 = extractConvPairParams(512, 512, 'mobilenetv1/conv_9')
const conv_10 = extractConvPairParams(512, 512, 'mobilenetv1/conv_10')
const conv_11 = extractConvPairParams(512, 512, 'mobilenetv1/conv_11')
const conv_12 = extractConvPairParams(512, 1024, 'mobilenetv1/conv_12')
const conv_13 = extractConvPairParams(1024, 1024, 'mobilenetv1/conv_13')
return {
conv_0_params,
conv_pair_params
conv_0,
conv_1,
conv_2,
conv_3,
conv_4,
conv_5,
conv_6,
conv_7,
conv_8,
conv_9,
conv_10,
conv_11,
conv_12,
conv_13
}
}
function extractPredictionLayerParams(): PredictionLayerParams {
const conv_0_params = extractPointwiseConvParams(1024, 256, 1)
const conv_1_params = extractPointwiseConvParams(256, 512, 3)
const conv_2_params = extractPointwiseConvParams(512, 128, 1)
const conv_3_params = extractPointwiseConvParams(128, 256, 3)
const conv_4_params = extractPointwiseConvParams(256, 128, 1)
const conv_5_params = extractPointwiseConvParams(128, 256, 3)
const conv_6_params = extractPointwiseConvParams(256, 64, 1)
const conv_7_params = extractPointwiseConvParams(64, 128, 3)
const box_encoding_0_predictor_params = extractConvParams(512, 12, 1)
const class_predictor_0_params = extractConvParams(512, 9, 1)
const box_encoding_1_predictor_params = extractConvParams(1024, 24, 1)
const class_predictor_1_params = extractConvParams(1024, 18, 1)
const box_encoding_2_predictor_params = extractConvParams(512, 24, 1)
const class_predictor_2_params = extractConvParams(512, 18, 1)
const box_encoding_3_predictor_params = extractConvParams(256, 24, 1)
const class_predictor_3_params = extractConvParams(256, 18, 1)
const box_encoding_4_predictor_params = extractConvParams(256, 24, 1)
const class_predictor_4_params = extractConvParams(256, 18, 1)
const box_encoding_5_predictor_params = extractConvParams(128, 24, 1)
const class_predictor_5_params = extractConvParams(128, 18, 1)
const box_predictor_0_params = {
box_encoding_predictor_params: box_encoding_0_predictor_params,
class_predictor_params: class_predictor_0_params
}
const box_predictor_1_params = {
box_encoding_predictor_params: box_encoding_1_predictor_params,
class_predictor_params: class_predictor_1_params
}
const box_predictor_2_params = {
box_encoding_predictor_params: box_encoding_2_predictor_params,
class_predictor_params: class_predictor_2_params
}
const box_predictor_3_params = {
box_encoding_predictor_params: box_encoding_3_predictor_params,
class_predictor_params: class_predictor_3_params
}
const box_predictor_4_params = {
box_encoding_predictor_params: box_encoding_4_predictor_params,
class_predictor_params: class_predictor_4_params
}
const box_predictor_5_params = {
box_encoding_predictor_params: box_encoding_5_predictor_params,
class_predictor_params: class_predictor_5_params
const conv_0 = extractPointwiseConvParams(1024, 256, 1, 'prediction_layer/conv_0')
const conv_1 = extractPointwiseConvParams(256, 512, 3, 'prediction_layer/conv_1')
const conv_2 = extractPointwiseConvParams(512, 128, 1, 'prediction_layer/conv_2')
const conv_3 = extractPointwiseConvParams(128, 256, 3, 'prediction_layer/conv_3')
const conv_4 = extractPointwiseConvParams(256, 128, 1, 'prediction_layer/conv_4')
const conv_5 = extractPointwiseConvParams(128, 256, 3, 'prediction_layer/conv_5')
const conv_6 = extractPointwiseConvParams(256, 64, 1, 'prediction_layer/conv_6')
const conv_7 = extractPointwiseConvParams(64, 128, 3, 'prediction_layer/conv_7')
const box_encoding_0_predictor = extractConvParams(512, 12, 1, 'prediction_layer/box_predictor_0/box_encoding_predictor')
const class_predictor_0 = extractConvParams(512, 9, 1, 'prediction_layer/box_predictor_0/class_predictor')
const box_encoding_1_predictor = extractConvParams(1024, 24, 1, 'prediction_layer/box_predictor_1/box_encoding_predictor')
const class_predictor_1 = extractConvParams(1024, 18, 1, 'prediction_layer/box_predictor_1/class_predictor')
const box_encoding_2_predictor = extractConvParams(512, 24, 1, 'prediction_layer/box_predictor_2/box_encoding_predictor')
const class_predictor_2 = extractConvParams(512, 18, 1, 'prediction_layer/box_predictor_2/class_predictor')
const box_encoding_3_predictor = extractConvParams(256, 24, 1, 'prediction_layer/box_predictor_3/box_encoding_predictor')
const class_predictor_3 = extractConvParams(256, 18, 1, 'prediction_layer/box_predictor_3/class_predictor')
const box_encoding_4_predictor = extractConvParams(256, 24, 1, 'prediction_layer/box_predictor_4/box_encoding_predictor')
const class_predictor_4 = extractConvParams(256, 18, 1, 'prediction_layer/box_predictor_4/class_predictor')
const box_encoding_5_predictor = extractConvParams(128, 24, 1, 'prediction_layer/box_predictor_5/box_encoding_predictor')
const class_predictor_5 = extractConvParams(128, 18, 1, 'prediction_layer/box_predictor_5/class_predictor')
const box_predictor_0 = {
box_encoding_predictor: box_encoding_0_predictor,
class_predictor: class_predictor_0
}
const box_predictor_1 = {
box_encoding_predictor: box_encoding_1_predictor,
class_predictor: class_predictor_1
}
const box_predictor_2 = {
box_encoding_predictor: box_encoding_2_predictor,
class_predictor: class_predictor_2
}
const box_predictor_3 = {
box_encoding_predictor: box_encoding_3_predictor,
class_predictor: class_predictor_3
}
const box_predictor_4 = {
box_encoding_predictor: box_encoding_4_predictor,
class_predictor: class_predictor_4
}
const box_predictor_5 = {
box_encoding_predictor: box_encoding_5_predictor,
class_predictor: class_predictor_5
}
return {
conv_0_params,
conv_1_params,
conv_2_params,
conv_3_params,
conv_4_params,
conv_5_params,
conv_6_params,
conv_7_params,
box_predictor_0_params,
box_predictor_1_params,
box_predictor_2_params,
box_predictor_3_params,
box_predictor_4_params,
box_predictor_5_params
conv_0,
conv_1,
conv_2,
conv_3,
conv_4,
conv_5,
conv_6,
conv_7,
box_predictor_0,
box_predictor_1,
box_predictor_2,
box_predictor_3,
box_predictor_4,
box_predictor_5
}
}
return {
extractMobilenetV1Params,
extractPredictionLayerParams
......@@ -173,7 +192,10 @@ function extractorsFactory(extractWeights: (numWeights: number) => Float32Array)
}
export function extractParams(weights: Float32Array): NetParams {
export function extractParams(weights: Float32Array): { params: NetParams, paramMappings: ParamMapping[] } {
const paramMappings: ParamMapping[] = []
const {
extractWeights,
getRemainingWeights
......@@ -182,25 +204,30 @@ export function extractParams(weights: Float32Array): NetParams {
const {
extractMobilenetV1Params,
extractPredictionLayerParams
} = extractorsFactory(extractWeights)
} = extractorsFactory(extractWeights, paramMappings)
const mobilenetv1_params = extractMobilenetV1Params()
const prediction_layer_params = extractPredictionLayerParams()
const mobilenetv1 = extractMobilenetV1Params()
const prediction_layer = extractPredictionLayerParams()
const extra_dim = tf.tensor3d(
extractWeights(5118 * 4),
[1, 5118, 4]
)
const output_layer_params = {
const output_layer = {
extra_dim
}
paramMappings.push({ paramPath: 'output_layer/extra_dim' })
if (getRemainingWeights().length !== 0) {
throw new Error(`weights remaing after extract: ${getRemainingWeights().length}`)
}
return {
mobilenetv1_params,
prediction_layer_params,
output_layer_params
params: {
mobilenetv1,
prediction_layer,
output_layer
},
paramMappings
}
}
\ No newline at end of file
src/faceDetectionNet/loadQuantizedParams.ts
import { isTensor1D, isTensor4D, isTensor3D } from '../commons/isTensor';
import { tf } from '..';
import { disposeUnusedWeightTensors } from '../commons/disposeUnusedWeightTensors';
import { extractWeightEntryFactory } from '../commons/extractWeightEntryFactory';
import { isTensor1D, isTensor3D, isTensor4D } from '../commons/isTensor';
import { loadWeightMap } from '../commons/loadWeightMap';
import { BoxPredictionParams, MobileNetV1, PointwiseConvParams, PredictionLayerParams } from './types';
import { ConvParams, ParamMapping } from '../commons/types';
import { BoxPredictionParams, MobileNetV1, NetParams, PointwiseConvParams, PredictionLayerParams } from './types';
const DEFAULT_MODEL_NAME = 'face_detection_model'
function extractorsFactory(weightMap: any) {
function extractorsFactory(weightMap: any, paramMappings: ParamMapping[]) {
function extractPointwiseConvParams(prefix: string, idx: number): PointwiseConvParams {
const extractWeightEntry = extractWeightEntryFactory(weightMap, paramMappings)
const pointwise_conv_params = {
filters: weightMap[`${prefix}/Conv2d_${idx}_pointwise/weights`],
batch_norm_offset: weightMap[`${prefix}/Conv2d_${idx}_pointwise/convolution_bn_offset`]
}
if (!isTensor4D(pointwise_conv_params.filters)) {
throw new Error(`expected weightMap[${prefix}/Conv2d_${idx}_pointwise/weights] to be a Tensor4D, instead have ${pointwise_conv_params.filters}`)
}
function extractPointwiseConvParams(prefix: string, idx: number, mappedPrefix: string): PointwiseConvParams {
if (!isTensor1D(pointwise_conv_params.batch_norm_offset)) {
throw new Error(`expected weightMap[${prefix}/Conv2d_${idx}_pointwise/convolution_bn_offset] to be a Tensor1D, instead have ${pointwise_conv_params.batch_norm_offset}`)
}
const filters = extractWeightEntry<tf.Tensor4D>(`${prefix}/Conv2d_${idx}_pointwise/weights`, 4, `${mappedPrefix}/filters`)
const batch_norm_offset = extractWeightEntry<tf.Tensor1D>(`${prefix}/Conv2d_${idx}_pointwise/convolution_bn_offset`, 1, `${mappedPrefix}/batch_norm_offset`)
return pointwise_conv_params
return { filters, batch_norm_offset }
}
function extractConvPairParams(idx: number): MobileNetV1.ConvPairParams {
const depthwise_conv_params = {
filters: weightMap[`MobilenetV1/Conv2d_${idx}_depthwise/depthwise_weights`],
batch_norm_scale: weightMap[`MobilenetV1/Conv2d_${idx}_depthwise/BatchNorm/gamma`],
batch_norm_offset: weightMap[`MobilenetV1/Conv2d_${idx}_depthwise/BatchNorm/beta`],
batch_norm_mean: weightMap[`MobilenetV1/Conv2d_${idx}_depthwise/BatchNorm/moving_mean`],
batch_norm_variance: weightMap[`MobilenetV1/Conv2d_${idx}_depthwise/BatchNorm/moving_variance`],
}
if (!isTensor4D(depthwise_conv_params.filters)) {
throw new Error(`expected weightMap[MobilenetV1/Conv2d_${idx}_depthwise/depthwise_weights] to be a Tensor4D, instead have ${depthwise_conv_params.filters}`)
}
if (!isTensor1D(depthwise_conv_params.batch_norm_scale)) {
throw new Error(`expected weightMap[MobilenetV1/Conv2d_${idx}_depthwise/BatchNorm/gamma] to be a Tensor1D, instead have ${depthwise_conv_params.batch_norm_scale}`)
}
if (!isTensor1D(depthwise_conv_params.batch_norm_offset)) {
throw new Error(`expected weightMap[MobilenetV1/Conv2d_${idx}_depthwise/BatchNorm/beta] to be a Tensor1D, instead have ${depthwise_conv_params.batch_norm_offset}`)
}
if (!isTensor1D(depthwise_conv_params.batch_norm_mean)) {
throw new Error(`expected weightMap[MobilenetV1/Conv2d_${idx}_depthwise/BatchNorm/moving_mean] to be a Tensor1D, instead have ${depthwise_conv_params.batch_norm_mean}`)
}
const mappedPrefix = `mobilenetv1/conv_${idx}`
const prefixDepthwiseConv = `MobilenetV1/Conv2d_${idx}_depthwise`
const mappedPrefixDepthwiseConv = `${mappedPrefix}/depthwise_conv`
const mappedPrefixPointwiseConv = `${mappedPrefix}/pointwise_conv`
if (!isTensor1D(depthwise_conv_params.batch_norm_variance)) {
throw new Error(`expected weightMap[MobilenetV1/Conv2d_${idx}_depthwise/BatchNorm/moving_variance] to be a Tensor1D, instead have ${depthwise_conv_params.batch_norm_variance}`)
}
const filters = extractWeightEntry<tf.Tensor4D>(`${prefixDepthwiseConv}/depthwise_weights`, 4, `${mappedPrefixDepthwiseConv}/filters`)
const batch_norm_scale = extractWeightEntry<tf.Tensor1D>(`${prefixDepthwiseConv}/BatchNorm/gamma`, 1, `${mappedPrefixDepthwiseConv}/batch_norm_scale`)
const batch_norm_offset = extractWeightEntry<tf.Tensor1D>(`${prefixDepthwiseConv}/BatchNorm/beta`, 1, `${mappedPrefixDepthwiseConv}/batch_norm_offset`)
const batch_norm_mean = extractWeightEntry<tf.Tensor1D>(`${prefixDepthwiseConv}/BatchNorm/moving_mean`, 1, `${mappedPrefixDepthwiseConv}/batch_norm_mean`)
const batch_norm_variance = extractWeightEntry<tf.Tensor1D>(`${prefixDepthwiseConv}/BatchNorm/moving_variance`, 1, `${mappedPrefixDepthwiseConv}/batch_norm_variance`)
return {
depthwise_conv_params,
pointwise_conv_params: extractPointwiseConvParams('MobilenetV1', idx)
depthwise_conv: {
filters,
batch_norm_scale,
batch_norm_offset,
batch_norm_mean,
batch_norm_variance
},
pointwise_conv: extractPointwiseConvParams('MobilenetV1', idx, mappedPrefixPointwiseConv)
}
}
function extractMobilenetV1Params(): MobileNetV1.Params {
return {
conv_0_params: extractPointwiseConvParams('MobilenetV1', 0),
conv_pair_params: Array(13).fill(0).map((_, i) => extractConvPairParams(i + 1))
}
}
function extractBoxPredictorParams(idx: number): BoxPredictionParams {
const params = {
box_encoding_predictor_params: {
filters: weightMap[`Prediction/BoxPredictor_${idx}/BoxEncodingPredictor/weights`],
bias: weightMap[`Prediction/BoxPredictor_${idx}/BoxEncodingPredictor/biases`]
},
class_predictor_params: {
filters: weightMap[`Prediction/BoxPredictor_${idx}/ClassPredictor/weights`],
bias: weightMap[`Prediction/BoxPredictor_${idx}/ClassPredictor/biases`]
conv_0: extractPointwiseConvParams('MobilenetV1', 0, 'mobilenetv1/conv_0'),
conv_1: extractConvPairParams(1),
conv_2: extractConvPairParams(2),
conv_3: extractConvPairParams(3),
conv_4: extractConvPairParams(4),
conv_5: extractConvPairParams(5),
conv_6: extractConvPairParams(6),
conv_7: extractConvPairParams(7),
conv_8: extractConvPairParams(8),
conv_9: extractConvPairParams(9),
conv_10: extractConvPairParams(10),
conv_11: extractConvPairParams(11),
conv_12: extractConvPairParams(12),
conv_13: extractConvPairParams(13)
}
}
if (!isTensor4D(params.box_encoding_predictor_params.filters)) {
throw new Error(`expected weightMap[Prediction/BoxPredictor_${idx}/BoxEncodingPredictor/weights] to be a Tensor4D, instead have ${params.box_encoding_predictor_params.filters}`)
}
function extractConvParams(prefix: string, mappedPrefix: string): ConvParams {
const filters = extractWeightEntry<tf.Tensor4D>(`${prefix}/weights`, 4, `${mappedPrefix}/filters`)
const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/biases`, 1, `${mappedPrefix}/bias`)
if (!isTensor1D(params.box_encoding_predictor_params.bias)) {
throw new Error(`expected weightMap[Prediction/BoxPredictor_${idx}/BoxEncodingPredictor/biases] to be a Tensor1D, instead have ${params.box_encoding_predictor_params.bias}`)
return { filters, bias }
}
if (!isTensor4D(params.class_predictor_params.filters)) {
throw new Error(`expected weightMap[Prediction/BoxPredictor_${idx}/ClassPredictor/weights] to be a Tensor4D, instead have ${params.class_predictor_params.filters}`)
}
function extractBoxPredictorParams(idx: number): BoxPredictionParams {
if (!isTensor1D(params.class_predictor_params.bias)) {
throw new Error(`expected weightMap[Prediction/BoxPredictor_${idx}/ClassPredictor/biases] to be a Tensor1D, instead have ${params.class_predictor_params.bias}`)
}
const box_encoding_predictor = extractConvParams(
`Prediction/BoxPredictor_${idx}/BoxEncodingPredictor`,
`prediction_layer/box_predictor_${idx}/box_encoding_predictor`
)
const class_predictor = extractConvParams(
`Prediction/BoxPredictor_${idx}/ClassPredictor`,
`prediction_layer/box_predictor_${idx}/class_predictor`
)
return params
return { box_encoding_predictor, class_predictor }
}
function extractPredictionLayerParams(): PredictionLayerParams {
return {
conv_0_params: extractPointwiseConvParams('Prediction', 0),
conv_1_params: extractPointwiseConvParams('Prediction', 1),
conv_2_params: extractPointwiseConvParams('Prediction', 2),
conv_3_params: extractPointwiseConvParams('Prediction', 3),
conv_4_params: extractPointwiseConvParams('Prediction', 4),
conv_5_params: extractPointwiseConvParams('Prediction', 5),
conv_6_params: extractPointwiseConvParams('Prediction', 6),
conv_7_params: extractPointwiseConvParams('Prediction', 7),
box_predictor_0_params: extractBoxPredictorParams(0),
box_predictor_1_params: extractBoxPredictorParams(1),
box_predictor_2_params: extractBoxPredictorParams(2),
box_predictor_3_params: extractBoxPredictorParams(3),
box_predictor_4_params: extractBoxPredictorParams(4),
box_predictor_5_params: extractBoxPredictorParams(5)
conv_0: extractPointwiseConvParams('Prediction', 0, 'prediction_layer/conv_0'),
conv_1: extractPointwiseConvParams('Prediction', 1, 'prediction_layer/conv_1'),
conv_2: extractPointwiseConvParams('Prediction', 2, 'prediction_layer/conv_2'),
conv_3: extractPointwiseConvParams('Prediction', 3, 'prediction_layer/conv_3'),
conv_4: extractPointwiseConvParams('Prediction', 4, 'prediction_layer/conv_4'),
conv_5: extractPointwiseConvParams('Prediction', 5, 'prediction_layer/conv_5'),
conv_6: extractPointwiseConvParams('Prediction', 6, 'prediction_layer/conv_6'),
conv_7: extractPointwiseConvParams('Prediction', 7, 'prediction_layer/conv_7'),
box_predictor_0: extractBoxPredictorParams(0),
box_predictor_1: extractBoxPredictorParams(1),
box_predictor_2: extractBoxPredictorParams(2),
box_predictor_3: extractBoxPredictorParams(3),
box_predictor_4: extractBoxPredictorParams(4),
box_predictor_5: extractBoxPredictorParams(5)
}
}
......@@ -124,24 +110,34 @@ function extractorsFactory(weightMap: any) {
}
}
export async function loadQuantizedParams(uri: string | undefined): Promise<any> {//Promise<NetParams> {
export async function loadQuantizedParams(
uri: string | undefined
): Promise<{ params: NetParams, paramMappings: ParamMapping[] }> {
const weightMap = await loadWeightMap(uri, DEFAULT_MODEL_NAME)
const paramMappings: ParamMapping[] = []
const {
extractMobilenetV1Params,
extractPredictionLayerParams
} = extractorsFactory(weightMap)
} = extractorsFactory(weightMap, paramMappings)
const extra_dim = weightMap['Output/extra_dim']
paramMappings.push({ originalPath: 'Output/extra_dim', paramPath: 'output_layer/extra_dim' })
if (!isTensor3D(extra_dim)) {
throw new Error(`expected weightMap['Output/extra_dim'] to be a Tensor3D, instead have ${extra_dim}`)
}
return {
mobilenetv1_params: extractMobilenetV1Params(),
prediction_layer_params: extractPredictionLayerParams(),
output_layer_params: {
const params = {
mobilenetv1: extractMobilenetV1Params(),
prediction_layer: extractPredictionLayerParams(),
output_layer: {
extra_dim
}
}
disposeUnusedWeightTensors(weightMap, paramMappings)
return { params, paramMappings }
}
\ No newline at end of file
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
constructor() {
super('FaceRecognitionNet')
}
if (weightsOrUrl && typeof weightsOrUrl !== 'string') {
throw new Error('FaceLandmarkNet.load - expected model uri, or weights as Float32Array')
}
this._params = await loadQuantizedParams(weightsOrUrl)
}
public forwardInput(input: NetInput): tf.Tensor2D {
public extractWeights(weights: Float32Array) {
this._params = extractParams(weights)
}
const { params } = this
public forwardInput(input: NetInput): tf.Tensor2D {
if (!this._params) {
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(
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}`)
}
const filters = extractWeightEntry<tf.Tensor4D>(`${prefix}/conv/filters`, 4)
const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/conv/bias`, 1)
const scale = extractScaleLayerParams(prefix)
if (!isTensor1D(params.bias)) {
throw new Error(`expected weightMap[${prefix}/conv/bias] to be a Tensor1D, instead have ${params.bias}`)
}
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
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment