TensorFlow.js自定义含子层无自有权重层训练时报无训练变量错误

问题原因 #

TensorFlow.js 与 Python Keras 在子层管理逻辑上存在差异：Python Keras 会自动追踪并注册构造函数中创建的子层，而 TensorFlow.js 的自定义层不会自动识别子层。如果不在父层中显式注册子层，父层无法收集子层的可训练变量，导致model.fit时抛出variableGrads() expects at least one of the input variables to be trainable错误。

解决方案 #

在自定义层的构造函数中创建子层后，必须调用this.addLayer(子层实例)完成注册，让父层递归管理子层的可训练权重。

修正后的代码 #

import * as tf from "@tensorflow/tfjs";
import { LayerArgs } from "@tensorflow/tfjs-layers/dist/engine/topology";
import { Kwargs } from "@tensorflow/tfjs-layers/dist/types";

class DenseCustom extends tf.layers.Layer {
    private static _className = "DenseCustom";

    public dense: tf.layers.Layer;

    constructor(args: LayerArgs & { units: number }) {
        super(args);

        this.dense = tf.layers.dense({ units: args.units, useBias: false });
        // 关键步骤：将子层注册到父层，确保权重被追踪
        this.addLayer(this.dense);
    }

    override call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[] {
        return this.dense.apply(inputs) as tf.Tensor;
    }

    // 实现getConfig以支持模型序列化（可选但推荐）
    override getConfig(): tf.serialization.ConfigDict {
        const baseConfig = super.getConfig();
        return { ...baseConfig, units: this.dense.getConfig().units };
    }

    // 自定义层必须暴露className用于序列化
    static get className() {
        return DenseCustom._className;
    }
}

// 注册自定义层到TensorFlow.js的序列化系统
tf.serialization.registerClass(DenseCustom);

补充说明 #

1. addLayer()方法会将子层加入父层的管理队列，父层会自动合并子层的trainableWeights和nonTrainableWeights，确保训练时能识别到可训练变量。
2. 实现getConfig和注册className是为了支持模型的保存与加载，若不需要序列化可省略，但生产环境建议添加。

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