import * as tf from '@tensorflow/tfjs';
import { InputSpec } from '@tensorflow/tfjs-layers/dist/engine/topology';
class TT extends tf.layers.Layer {
public static className = 'TT';
public className = TT.className;
public timed: tf.layers.Layer | null = null;
constructor(config?: any) {
super(config || {});
this.inputSpec = [new InputSpec({ ndim: 3 })];
}
public build(inputShape: tf.Shape): void {
const dimensions = inputShape[2];
const timed = tf.sequential({ name: 'per_time_step' });
timed.add(tf.layers.dense({ inputShape: [dimensions], kernelInitializer: 'zeros', units: dimensions }));
timed.add(tf.layers.activation({ activation: 'softmax' }));
this.timed = tf.layers.timeDistributed({ layer: timed });
this.trainableWeights = this.timed.trainableWeights;
this.nonTrainableWeights = this.timed.nonTrainableWeights;
this.built = true;
}
public call(inputs: tf.Tensor[], kwargs: any) {
if (!this.built || !this.timed) { throw new Error(); }
return tf.tidy(() => {
this.invokeCallHook(inputs, kwargs);
return (this.timed as tf.layers.Layer).apply(inputs) as tf.Tensor;
});
}
public computeOutputShape(inputShape: tf.Shape) { return inputShape; }
}
tf.serialization.SerializationMap.register(TT);
const inputs = tf.input({ dtype: 'float32', shape: [1, 2] });
const lstm = tf.layers.lstm({ units: 2, returnSequences: true }).apply(inputs) as tf.SymbolicTensor;
const timeAttention = new TT({}).apply(lstm) as tf.SymbolicTensor;
const model = tf.model({ inputs, outputs: timeAttention });
const optimize = tf.train.adam(0.0066, 0.0025, 0.1);
model.compile({ loss: 'categoricalCrossentropy', metrics: ['accuracy'], optimizer: optimize });
const inp = tf.tensor3d([[[1,1]],[[2,2]],[[3,3]],[[4,4]],[[5,5]],[[6,6]]],[6,1,2]);
const out = tf.tensor3d([[[1,0]],[[2,0]],[[3,0]],[[4,0]],[[5,0]],[[6,0]]], [6, 1, 2]);
(async () => {
// NOTE: Here is the issue, if validationSplit is 0.5 this code works, else it fails
await model.fit(inp, out, { validationSplit: 0.5 });
})();