first implementation of TensorArray and preprocessor layer

justadudewhohacks · justadudewhohacks · commit 42e41f1e325c · 2018-06-02T15:22:46.000+02:00
diff --git a/src/faceDetectionNet/preprocessor/index.ts b/src/faceDetectionNet/preprocessor/index.ts
@@ -0,0 +1,29 @@
+import * as tf from '@tensorflow/tfjs-core';
+import { TensorArray } from '../../tfcpatches/TensorArray';
+import { whileLayer } from './whileLayer';
+
+export type PreprocessorParams = any
+
+// TODO: hardcoded params
+const elementShape = [512, 512, 3]
+const weight = tf.scalar(0.007843137718737125)
+const bias = tf.scalar(1)
+
+export function preprocessor(imgTensor: tf.Tensor4D, params: PreprocessorParams) {
+  const batchSize = imgTensor.shape[0]
+  const batchSizedArray1 = new TensorArray(batchSize, 'float32')
+  const batchSizedArray2 = new TensorArray(batchSize, 'float32')
+
+  let unusedFlow = null
+  const indices = tf.range(0, batchSize, 1)
+
+  // unstack
+  unusedFlow = batchSizedArray1.scatter(indices, imgTensor, unusedFlow)
+
+  unusedFlow = whileLayer(batchSizedArray1, batchSizedArray2, batchSize, unusedFlow)
+
+  // stack
+  const stacked = batchSizedArray2.gather(indices, unusedFlow, 'float32', elementShape)
+
+  return tf.add(tf.mul(stacked, weight), bias)
+}
diff --git a/src/faceDetectionNet/preprocessor/whileLayer.ts b/src/faceDetectionNet/preprocessor/whileLayer.ts
@@ -0,0 +1,9 @@
+import * as tf from '@tensorflow/tfjs-core';
+import { TensorArray } from '../../tfcpatches/TensorArray';
+
+export function whileLayer(arr1: TensorArray, arr2: TensorArray, batchSize: number, unusedFlowIn: tf.Scalar): tf.Scalar {
+  // TODO
+
+  const unusedFlowOut = tf.scalar(0)
+  return unusedFlowOut
+}
diff --git a/src/tfcpatches/TensorArray.ts b/src/tfcpatches/TensorArray.ts
@@ -0,0 +1,48 @@
+import * as tf from '@tensorflow/tfjs-core';
+
+export class TensorArray {
+  private _tensors: tf.Tensor[] | undefined
+
+  constructor(
+    private _size: number,
+    private _dtype: tf.DataType = null,
+    private _elementShape: number[] = null,
+    private _dynamicSize: boolean = false,
+    private _clearAfterRead: boolean = true,
+    private _identicalElementShapes: boolean = false,
+    private _tensorArrayName: string = null
+  ) {
+    if (_size) {
+      this._tensors = Array(_size).fill(0).map(_ => tf.scalar(0))
+    }
+  }
+
+  public scatter(indices: tf.Tensor1D, value: tf.Tensor, unusedFlow: tf.Scalar): tf.Scalar {
+    if (indices.shape.length !== 1) {
+      throw new Error(`scatter - expected rank of indices (${indices.shape.length}) to be 1`)
+    }
+    if (indices.shape[0] > this._size) {
+      throw new Error(`scatter - expected indices.shape[0] (${indices.shape[0]}) to be >= this._size (${this.size})`)
+    }
+    if (indices.shape[0] !== value.shape[0]) {
+      throw new Error(`scatter - expected indices.shape[0] (${indices.shape[0]}) to equal value.shape[0] (${value.shape[0]})`)
+    }
+
+    const unstacked = tf.unstack(value, 0)
+    Array.from(indices.dataSync()).forEach((idx, i) => {
+      this._tensors[idx] = unstacked[i]
+    })
+
+    const unusedFlowOut = tf.scalar(0)
+    return unusedFlowOut
+  }
+
+  public gather(indices: tf.Tensor1D, unusedFlow: tf.Scalar, dtype?: tf.DataType, elementShape?: number[]) : tf.Tensor {
+    const tensors = Array.from(indices.dataSync()).map(idx => this._tensors[idx])
+    return tf.concat(tensors)
+  }
+
+  public size(unusedFlow: tf.Scalar) {
+    return this._size
+  }
+}
