quantized weights + loading for tiny yolov2

src/tinyYolov2/TinyYolov2.ts

@@ -12,6 +12,7 @@ import { BOX_ANCHORS, INPUT_SIZES, IOU_THRESHOLD, NUM_BOXES } from './config';
 import { convWithBatchNorm } from './convWithBatchNorm';
 import { extractParams } from './extractParams';
 import { getDefaultParams } from './getDefaultParams';
+import { loadQuantizedParams } from './loadQuantizedParams';
 import { NetParams, TinyYolov2ForwardParams } from './types';
 
 export class TinyYolov2 extends NeuralNetwork<NetParams> {
@@ -135,11 +136,9 @@ export class TinyYolov2 extends NeuralNetwork<NetParams> {
     return detections
   }
 
-  /* TODO
   protected loadQuantizedParams(uri: string | undefined) {
     return loadQuantizedParams(uri)
   }
-  */
 
   protected extractParams(weights: Float32Array) {
     return extractParams(weights)

src/tinyYolov2/loadQuantizedParams.ts

+import * as tf from '@tensorflow/tfjs-core';
+
+import { disposeUnusedWeightTensors } from '../commons/disposeUnusedWeightTensors';
+import { extractWeightEntryFactory } from '../commons/extractWeightEntryFactory';
+import { loadWeightMap } from '../commons/loadWeightMap';
+import { ConvParams, ParamMapping } from '../commons/types';
+import { BatchNorm, ConvWithBatchNorm, NetParams } from './types';
+
+const DEFAULT_MODEL_NAME = 'tiny_yolov2_model'
+
+function extractorsFactory(weightMap: any, paramMappings: ParamMapping[]) {
+
+  const extractWeightEntry = extractWeightEntryFactory(weightMap, paramMappings)
+
+  function extractBatchNormParams(prefix: string): BatchNorm {
+    const sub = extractWeightEntry<tf.Tensor1D>(`${prefix}/sub`, 1)
+    const truediv = extractWeightEntry<tf.Tensor1D>(`${prefix}/truediv`, 1)
+    return { sub, truediv }
+  }
+
+  function extractConvParams(prefix: string): ConvParams {
+    const filters = extractWeightEntry<tf.Tensor4D>(`${prefix}/filters`, 4)
+    const bias = extractWeightEntry<tf.Tensor1D>(`${prefix}/bias`, 1)
+    return { filters, bias }
+  }
+
+  function extractConvWithBatchNormParams(prefix: string): ConvWithBatchNorm {
+    const conv = extractConvParams(`${prefix}/conv`)
+    const bn = extractBatchNormParams(`${prefix}/bn`)
+    return { conv, bn }
+  }
+
+  return {
+    extractConvParams,
+    extractConvWithBatchNormParams
+  }
+
+}
+
+export async function loadQuantizedParams(
+  uri: string | undefined
+): Promise<{ params: NetParams, paramMappings: ParamMapping[] }> {
+
+  const weightMap = await loadWeightMap(uri, DEFAULT_MODEL_NAME)
+  const paramMappings: ParamMapping[] = []
+
+  const {
+    extractConvParams,
+    extractConvWithBatchNormParams
+  } = extractorsFactory(weightMap, paramMappings)
+
+
+  const params = {
+    conv0: extractConvWithBatchNormParams('conv0'),
+    conv1: extractConvWithBatchNormParams('conv1'),
+    conv2: extractConvWithBatchNormParams('conv2'),
+    conv3: extractConvWithBatchNormParams('conv3'),
+    conv4: extractConvWithBatchNormParams('conv4'),
+    conv5: extractConvWithBatchNormParams('conv5'),
+    conv6: extractConvWithBatchNormParams('conv6'),
+    conv7: extractConvWithBatchNormParams('conv7'),
+    conv8: extractConvParams('conv8')
+  }
+
+  disposeUnusedWeightTensors(weightMap, paramMappings)
+
+  return { params, paramMappings }
+}
\ No newline at end of file

weights/tiny_yolov2_model-weights_manifest.json

+[{"weights":[{"name":"conv0/conv/filters","shape":[3,3,3,16],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0036540280370151294,"min":-0.44213739247883066}},{"name":"conv0/conv/bias","shape":[16],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.002164303555208094,"min":-0.13418682042290184}},{"name":"conv0/bn/sub","shape":[16],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0028029036288167914,"min":-0.3139252064274806}},{"name":"conv0/bn/truediv","shape":[16],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.05010388411727606,"min":2.7867696285247803}},{"name":"conv1/conv/filters","shape":[3,3,16,32],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0017398431020624497,"min":-0.22617960326811845}},{"name":"conv1/conv/bias","shape":[32],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.001145745112615473,"min":-0.12488621727508656}},{"name":"conv1/bn/sub","shape":[32],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.00586619587505565,"min":-0.9503237317590153}},{"name":"conv1/bn/truediv","shape":[32],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.007056442896525065,"min":1.2661800384521484}},{"name":"conv2/conv/filters","shape":[3,3,32,64],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0010545693776186774,"min":-0.12971203344709734}},{"name":"conv2/conv/bias","shape":[64],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0005329197820495157,"min":-0.07034541123053607}},{"name":"conv2/bn/sub","shape":[64],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.005863590567719703,"min":-0.691903686990925}},{"name":"conv2/bn/truediv","shape":[64],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.006374212339812634,"min":1.4001796245574951}},{"name":"conv3/conv/filters","shape":[3,3,64,128],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0006935241175632851,"min":-0.08114232175490436}},{"name":"conv3/conv/bias","shape":[128],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0002854522025468303,"min":-0.04224692597693088}},{"name":"conv3/bn/sub","shape":[128],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.008516784275279325,"min":-1.1667994457132675}},{"name":"conv3/bn/truediv","shape":[128],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.004096750652088838,"min":1.1394747495651245}},{"name":"conv4/conv/filters","shape":[3,3,128,256],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0006537227946169236,"min":-0.06733344784554313}},{"name":"conv4/conv/bias","shape":[256],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.00018004616905076831,"min":-0.041770711219778246}},{"name":"conv4/bn/sub","shape":[256],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.011157989969440535,"min":-1.5509606057522343}},{"name":"conv4/bn/truediv","shape":[256],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.00421752836190018,"min":1.1967103481292725}},{"name":"conv5/conv/filters","shape":[3,3,256,512],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0008280676077393925,"min":-0.05299632689532112}},{"name":"conv5/conv/bias","shape":[512],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.00007708504312105623,"min":-0.020973851904273033}},{"name":"conv5/bn/sub","shape":[512],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.011612189049814262,"min":-1.7302161684223252}},{"name":"conv5/bn/truediv","shape":[512],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0034869993434232826,"min":1.2963157892227173}},{"name":"conv6/conv/filters","shape":[3,3,512,1024],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0003250038945207409,"min":-0.04030048292057187}},{"name":"conv6/conv/bias","shape":[1024],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.00009164303220114009,"min":-0.02338002622127533}},{"name":"conv6/bn/sub","shape":[1024],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0168181017333386,"min":-2.068626513200648}},{"name":"conv6/bn/truediv","shape":[1024],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.004455030198190726,"min":1.309550166130066}},{"name":"conv7/conv/filters","shape":[3,3,1024,1024],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.00026194834533859704,"min":-0.033267439858001825}},{"name":"conv7/conv/bias","shape":[1024],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0001491123554753322,"min":-0.011928988438026577}},{"name":"conv7/bn/sub","shape":[1024],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.009263983427309523,"min":-1.1765258952683093}},{"name":"conv7/bn/truediv","shape":[1024],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.01123507022857666,"min":0.9360886812210083}},{"name":"conv8/filters","shape":[1,1,1024,30],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.001922343233052422,"min":-0.2614386796951294}},{"name":"conv8/bias","shape":[30],"dtype":"float32","quantization":{"dtype":"uint8","scale":0.0009638834233377494,"min":-0.15904076485072866}}],"paths":["tiny_yolov2_model-shard1","tiny_yolov2_model-shard2","tiny_yolov2_model-shard3","tiny_yolov2_model-shard4"]}]
\ No newline at end of file