Commit 42e41f1e by vincent

first implementation of TensorArray and preprocessor layer

parent 0387fd20
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)
}
\ No newline at end of file
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
}
\ No newline at end of file
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
}
}
\ 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