Closed
Description
Describe the current behavior
I implemented a definition layer and model using TensorFlowJS, and encountered a problem during training. The code is as follows.When the code runs, it will report the following error.
throw new Error("Error in gradient for op ".concat(node.kernelName, ". The gradient of input ") +
^
Error: Error in gradient for op BatchMatMul. The gradient of input 'b' has shape '4,8,8', which does not match the shape of the input '8,8'
Describe the expected behavior
no error
my code
import * as tf from '@tensorflow/tfjs-node';
class MyLayer extends tf.layers.Layer {
constructor(units) {
super({});
this.units = units;
}
build(inputShape) {
this.w1 = this.addWeight(
'w1',
[this.units, this.units],
'float32',
tf.initializers.glorotNormal({}),
undefined,
true
);
super.build(inputShape);
}
call(inputs) {
const input = Array.isArray(inputs) ? inputs[0] : inputs;
return tf.matMul(input, this.w1.read());
}
computeOutputShape(inputShape) {
return [null, inputShape[inputShape.length - 2], this.dModel];
}
static get className() {
return 'MyLayer';
}
}
tf.serialization.registerClass(MyLayer);
const input = tf.input({shape: [4, 8]});
const layer1 = new MyLayer(8, 2)
const output = layer1.apply(input)
const model = tf.model({inputs: input, outputs: output});
model.compile({
optimizer: 'adam',
loss: tf.losses.softmaxCrossEntropy,
});
const _input = tf.ones([40000, 4, 8])
const _output = tf.ones([40000, 4, 8])
model.fit(_input, _output, {batchSize: 4}).then(()=>{
let x = tf.ones([1, 4, 8])
const y = model.predict(x)
});