Skip to content

Commit 6f8221a

Browse files
justadudewhohacksjustadudewhohacks
committed
extractors for mobilenetv1 params
1 parent cb52d62 commit 6f8221a

File tree

11 files changed

+238
-70
lines changed

11 files changed

+238
-70
lines changed

‎src/faceDetectionNet/extractParams.ts

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
3+
import { FaceDetectionNet } from './types';
4+
5+
function mobilenetV1WeightsExtractorsFactory(extractWeights: (numWeights: number) => Float32Array) {
6+
7+
function extractDepthwiseConvParams(numChannels: number): FaceDetectionNet.MobileNetV1.DepthwiseConvParams {
8+
const weights = tf.tensor4d(extractWeights(3 * 3 * numChannels), [3, 3, numChannels, 1])
9+
const batch_norm_gamma = tf.tensor1d(extractWeights(numChannels))
10+
const batch_norm_beta = tf.tensor1d(extractWeights(numChannels))
11+
const batch_norm_mean = tf.tensor1d(extractWeights(numChannels))
12+
const batch_norm_variance = tf.tensor1d(extractWeights(numChannels))
13+
14+
return {
15+
weights,
16+
batch_norm_gamma,
17+
batch_norm_beta,
18+
batch_norm_mean,
19+
batch_norm_variance
20+
}
21+
}
22+
23+
function extractPointwiseConvParams(channelsIn: number, channelsOut: number): FaceDetectionNet.MobileNetV1.PointwiseConvParams {
24+
const weights = tf.tensor4d(extractWeights(channelsIn * channelsOut), [1, 1, channelsIn, channelsOut])
25+
const batch_norm_offset = tf.tensor1d(extractWeights(channelsOut))
26+
27+
return {
28+
weights,
29+
batch_norm_offset
30+
}
31+
}
32+
33+
function extractConvPairParams(channelsIn: number, channelsOut: number): FaceDetectionNet.MobileNetV1.ConvPairParams {
34+
const depthwise_conv_params = extractDepthwiseConvParams(channelsIn)
35+
const pointwise_conv_params = extractPointwiseConvParams(channelsIn, channelsOut)
36+
37+
return {
38+
depthwise_conv_params,
39+
pointwise_conv_params
40+
}
41+
}
42+
43+
return {
44+
extractPointwiseConvParams,
45+
extractConvPairParams
46+
}
47+
48+
}
49+
50+
function extractorsFactory(extractWeights: (numWeights: number) => Float32Array) {
51+
52+
const {
53+
extractPointwiseConvParams,
54+
extractConvPairParams
55+
} = mobilenetV1WeightsExtractorsFactory(extractWeights)
56+
57+
58+
59+
function extractMobilenetV1Params(): FaceDetectionNet.MobileNetV1.Params {
60+
61+
const conv_0_params = {
62+
weights: tf.tensor4d(extractWeights(3 * 3 * 3 * 32), [3, 3, 3, 32]),
63+
batch_norm_offset: tf.tensor1d(extractWeights(32))
64+
65+
}
66+
67+
const channelNumPairs = [
68+
[32, 64],
69+
[64, 128],
70+
[128, 128],
71+
[128, 256],
72+
[256, 256],
73+
[256, 512],
74+
[512, 512],
75+
[512, 512],
76+
[512, 512],
77+
[512, 512],
78+
[512, 512],
79+
[512, 1024],
80+
[1024, 1024]
81+
]
82+
83+
const conv_pair_params = channelNumPairs.map(
84+
([channelsIn, channelsOut]) => extractConvPairParams(channelsIn, channelsOut)
85+
)
86+
87+
return {
88+
conv_0_params,
89+
conv_pair_params
90+
}
91+
92+
}
93+
94+
95+
return {
96+
extractMobilenetV1Params
97+
}
98+
99+
}
100+
101+
export function extractParams(weights: Float32Array): FaceDetectionNet.NetParams {
102+
const extractWeights = (numWeights: number): Float32Array => {
103+
console.log(numWeights)
104+
const ret = weights.slice(0, numWeights)
105+
weights = weights.slice(numWeights)
106+
return ret
107+
}
108+
109+
const {
110+
extractMobilenetV1Params
111+
} = extractorsFactory(extractWeights)
112+
113+
const mobilenetv1_params = extractMobilenetV1Params()
114+
115+
if (weights.length !== 0) {
116+
throw new Error(`weights remaing after extract: ${weights.length}`)
117+
}
118+
119+
return {
120+
mobilenetv1_params
121+
}
122+
}

‎src/faceDetectionNet/index.ts

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
import * as tf from '@tensorflow/tfjs-core';
22

33
import { resizeLayer } from './resizeLayer';
4+
import { extractParams } from './extractParams';
5+
import { mobileNetV1 } from './mobileNetV1';
46

5-
export function faceDetectionNet() {
7+
export function faceDetectionNet(weights: Float32Array) {
8+
const params = extractParams(weights)
69

710
async function forward(input: ImageData|ImageData[]) {
811

@@ -12,12 +15,11 @@ export function faceDetectionNet() {
1215

1316
const imgTensor = tf.cast(tf.concat(imgTensors, 0), 'float32')
1417

15-
const resized = resizeLayer(imgTensor)
18+
let out = resizeLayer(imgTensor) as tf.Tensor4D
1619

20+
out = mobileNetV1(out, params.mobilenetv1_params)
1721

18-
19-
20-
return resized
22+
return out
2123
}
2224

2325
return {

‎src/faceDetectionNet/mobileNetV1.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
3+
import { FaceDetectionNet } from './types';
4+
5+
export function mobileNetV1(x: tf.Tensor4D, params: FaceDetectionNet.MobileNetV1.Params) {
6+
return x
7+
}

‎src/faceDetectionNet/resizeLayer.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ const resizedImageSize = [512, 512] as [number, number]
55
const weight = tf.scalar(0.007843137718737125)
66
const bias = tf.scalar(1)
77

8-
export function resizeLayer(imgTensor: tf.Tensor4D) {
9-
const resizedImgs = tf.image.resizeBilinear(imgTensor, resizedImageSize, false)
10-
return tf.sub(tf.mul(resizedImgs, weight), bias)
8+
export function resizeLayer(x: tf.Tensor4D) {
9+
const resized = tf.image.resizeBilinear(x, resizedImageSize, false)
10+
return tf.sub(tf.mul(resized, weight), bias)
1111
}

‎src/faceDetectionNet/types.ts

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import * as tf from '@tensorflow/tfjs-core';
2+
3+
export namespace FaceDetectionNet {
4+
5+
export namespace MobileNetV1 {
6+
7+
export type DepthwiseConvParams = {
8+
weights: tf.Tensor4D // [3, 3, ch, 1]
9+
batch_norm_gamma: tf.Tensor1D
10+
batch_norm_beta: tf.Tensor1D
11+
batch_norm_mean: tf.Tensor1D
12+
batch_norm_variance: tf.Tensor1D
13+
}
14+
15+
export type PointwiseConvParams = {
16+
weights: tf.Tensor4D // [1, 1, ch_in, ch_out]
17+
batch_norm_offset: tf.Tensor1D
18+
}
19+
20+
export type ConvPairParams = {
21+
depthwise_conv_params: DepthwiseConvParams
22+
pointwise_conv_params: PointwiseConvParams
23+
}
24+
25+
export type Params = {
26+
conv_0_params: PointwiseConvParams
27+
conv_pair_params: ConvPairParams[]
28+
}
29+
30+
}
31+
32+
export type NetParams = {
33+
mobilenetv1_params: MobileNetV1.Params
34+
}
35+
36+
}

‎src/faceRecognitionNet/FaceRecognitionNetParams.ts

Lines changed: 0 additions & 23 deletions
This file was deleted.

‎src/faceRecognitionNet/convLayer.ts

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,12 @@
11
import * as tf from '@tensorflow/tfjs-core';
22

3-
import { scale, ScaleLayerParams } from './scaleLayer';
3+
import { scale } from './scaleLayer';
4+
import { FaceRecognitionNet } from './types';
45

5-
export type ConvParams = {
6-
filters: tf.Tensor4D
7-
biases: tf.Tensor1D
8-
}
9-
10-
export type ConvLayerParams = {
11-
conv: ConvParams
12-
scale: ScaleLayerParams
13-
}
146

157
function convLayer(
168
x: tf.Tensor4D,
17-
params: ConvLayerParams,
9+
params: FaceRecognitionNet.ConvLayerParams,
1810
stride: number,
1911
withRelu: boolean,
2012
padding: 'valid' | 'same' = 'same'
@@ -27,14 +19,14 @@ function convLayer(
2719
return withRelu ? tf.relu(out) : out
2820
}
2921

30-
export function conv(x: tf.Tensor4D, params: ConvLayerParams) {
22+
export function conv(x: tf.Tensor4D, params: FaceRecognitionNet.ConvLayerParams) {
3123
return convLayer(x, params, 1, true)
3224
}
3325

34-
export function convNoRelu(x: tf.Tensor4D, params: ConvLayerParams) {
26+
export function convNoRelu(x: tf.Tensor4D, params: FaceRecognitionNet.ConvLayerParams) {
3527
return convLayer(x, params, 1, false)
3628
}
3729

38-
export function convDown(x: tf.Tensor4D, params: ConvLayerParams) {
30+
export function convDown(x: tf.Tensor4D, params: FaceRecognitionNet.ConvLayerParams) {
3931
return convLayer(x, params, 2, true, 'valid')
4032
}

‎src/faceRecognitionNet/extractParams.ts

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,6 @@
11
import * as tf from '@tensorflow/tfjs-core';
22

3-
import { ConvLayerParams } from './convLayer';
4-
import { FaceRecognitionNetParams } from './FaceRecognitionNetParams';
5-
import { ResidualLayerParams } from './residualLayer';
6-
import { ScaleLayerParams } from './scaleLayer';
3+
import { FaceRecognitionNet } from './types';
74

85
function isFloat(num: number) {
96
return num % 1 !== 0
@@ -26,7 +23,7 @@ function extractorsFactory(extractWeights: (numWeights: number) => Float32Array)
2623
)
2724
}
2825

29-
function extractScaleLayerParams(numWeights: number): ScaleLayerParams {
26+
function extractScaleLayerParams(numWeights: number): FaceRecognitionNet.ScaleLayerParams {
3027
const weights = tf.tensor1d(extractWeights(numWeights))
3128
const biases = tf.tensor1d(extractWeights(numWeights))
3229
return {
@@ -35,7 +32,7 @@ function extractorsFactory(extractWeights: (numWeights: number) => Float32Array)
3532
}
3633
}
3734

38-
function extractConvLayerParams(numFilterValues: number, numFilters: number, filterSize: number): ConvLayerParams {
35+
function extractConvLayerParams(numFilterValues: number, numFilters: number, filterSize: number): FaceRecognitionNet.ConvLayerParams {
3936
const conv_filters = extractFilterValues(numFilterValues, numFilters, filterSize)
4037
const conv_biases = tf.tensor1d(extractWeights(numFilters))
4138
const scale = extractScaleLayerParams(numFilters)
@@ -49,9 +46,9 @@ function extractorsFactory(extractWeights: (numWeights: number) => Float32Array)
4946
}
5047
}
5148

52-
function extractResidualLayerParams(numFilterValues: number, numFilters: number, filterSize: number, isDown: boolean = false): ResidualLayerParams {
53-
const conv1: ConvLayerParams = extractConvLayerParams((isDown ? 0.5 : 1) * numFilterValues, numFilters, filterSize)
54-
const conv2: ConvLayerParams = extractConvLayerParams(numFilterValues, numFilters, filterSize)
49+
function extractResidualLayerParams(numFilterValues: number, numFilters: number, filterSize: number, isDown: boolean = false): FaceRecognitionNet.ResidualLayerParams {
50+
const conv1: FaceRecognitionNet.ConvLayerParams = extractConvLayerParams((isDown ? 0.5 : 1) * numFilterValues, numFilters, filterSize)
51+
const conv2: FaceRecognitionNet.ConvLayerParams = extractConvLayerParams(numFilterValues, numFilters, filterSize)
5552

5653
return {
5754
conv1,
@@ -66,7 +63,7 @@ function extractorsFactory(extractWeights: (numWeights: number) => Float32Array)
6663

6764
}
6865

69-
export function extractParams(weights: Float32Array): FaceRecognitionNetParams {
66+
export function extractParams(weights: Float32Array): FaceRecognitionNet.NetParams {
7067
const extractWeights = (numWeights: number): Float32Array => {
7168
const ret = weights.slice(0, numWeights)
7269
weights = weights.slice(numWeights)

‎src/faceRecognitionNet/residualLayer.ts

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,17 @@
11
import * as tf from '@tensorflow/tfjs-core';
22

3-
import { conv, convNoRelu, ConvLayerParams, convDown } from './convLayer';
3+
import { conv, convDown, convNoRelu } from './convLayer';
4+
import { FaceRecognitionNet } from './types';
45

5-
export type ResidualLayerParams = {
6-
conv1: ConvLayerParams
7-
conv2: ConvLayerParams
8-
}
9-
10-
export function residual(x: tf.Tensor4D, params: ResidualLayerParams): tf.Tensor4D {
6+
export function residual(x: tf.Tensor4D, params: FaceRecognitionNet.ResidualLayerParams): tf.Tensor4D {
117
let out = conv(x, params.conv1)
128
out = convNoRelu(out, params.conv2)
139
out = tf.add(out, x)
1410
out = tf.relu(out)
1511
return out
1612
}
1713

18-
export function residualDown(x: tf.Tensor4D, params: ResidualLayerParams): tf.Tensor4D {
14+
export function residualDown(x: tf.Tensor4D, params: FaceRecognitionNet.ResidualLayerParams): tf.Tensor4D {
1915
let out = convDown(x, params.conv1)
2016
out = convNoRelu(out, params.conv2)
2117

‎src/faceRecognitionNet/scaleLayer.ts

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,6 @@
11
import * as tf from '@tensorflow/tfjs-core';
2+
import { FaceRecognitionNet } from './types';
23

3-
export type ScaleLayerParams = {
4-
weights: tf.Tensor1D
5-
biases: tf.Tensor1D
6-
}
7-
8-
export function scale(x: tf.Tensor4D, params: ScaleLayerParams): tf.Tensor4D {
4+
export function scale(x: tf.Tensor4D, params: FaceRecognitionNet.ScaleLayerParams): tf.Tensor4D {
95
return tf.add(tf.mul(x, params.weights), params.biases)
10-
}
6+
}

0 commit comments

Comments
 (0)