TensorFlowJS自定义Layer训练报错：BatchMatMul梯度形状不匹配求助

报错的核心原因是自定义层的computeOutputShape方法返回了错误的输出形状，同时代码中存在几处小问题导致梯度计算时形状不匹配。

具体问题点 #

• computeOutputShape中使用了未定义的this.dModel，应该替换为this.units，否则框架无法正确推断输出张量形状，进而导致梯度计算流程中形状校验失败。
• 实例化MyLayer时传入了多余的参数2，构造函数仅接受units一个参数，该多余参数无意义。
• call方法中使用this.w1.read()并非必要，直接使用this.w1即可（TFJS自定义层的权重变量在call方法中会自动被梯度跟踪）。

修正后的完整代码 #

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;
    // 直接使用权重变量，无需read()
    return tf.matMul(input, this.w1);
  }

  computeOutputShape(inputShape) {
    // 替换未定义的this.dModel为this.units，正确推断输出形状
    return [null, inputShape[inputShape.length - 2], this.units];
  }

  static get className() {
    return 'MyLayer';
  }
}
tf.serialization.registerClass(MyLayer);

const input = tf.input({shape: [4, 8]});
// 仅传入units参数8，移除多余的2
const layer1 = new MyLayer(8);
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);
  y.print();
});

额外说明 #

• 关于损失函数：tf.losses.softmaxCrossEntropy要求标签是one-hot编码形式，当前用tf.ones生成的标签不符合该损失函数的预期，会导致训练时损失值异常，建议根据实际任务调整标签格式。
• 批量矩阵乘法：输入张量形状为[batch, 4, 8]，权重形状为[8,8]，tf.matMul会自动对每个batch内的[4,8]矩阵与权重做乘法，输出[batch,4,8]，这部分逻辑是正确的。

内容的提问来源于stack exchange，提问作者xinghuode