Standard layers (like `tf.layers.dense`) fail to report their trainable weights, when used from within a custom layer about tfjs HOT 4 OPEN

So, it does look like tf.layers.activation should work here, as a replacement to the string activations. I hadn't noticed this layer type before:
https://js.tensorflow.org/api/latest/#layers.activation

The other stuff is still a problem, so far as I can see.

from tfjs.

You haven't registered the weights from the child layers to the parent later In the build method.

this.trainableWeights = [...this.childLayer.trainableWeights]

Same goes for nonTrainableWeights

We really need some documentation written up for the layers API, especially the rnncell base class

from tfjs.

Thanks for the suggestion, @lukemovement. Your suggestion does work, however you also need to build each of the child layers:

this.inProj.build(inputShape)
this.outProj.build(inputShape)

I actually discovered this trick previously, and it also has problems. The issue relates to restoring a model from disk:

/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:273
        var _this = _super.call(this, message) || this;
                           ^

ValueError: Duplicate weight name: glu-TLd/kernel
    at new ValueError (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:273:28)
    at Container.loadWeights (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:21823:35)
    at /home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:25792:27
    at step (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:159:27)
    at Object.next (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:108:53)
    at fulfilled (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:89:28)

Node.js v18.18.2

When you add child layers to trainableWeights like this, every child layer will inherit the name of its parent (with a suffix):

  LayerVariable {
    dtype: 'float32',
    shape: [ 333 ],
    id: 30,
    originalName: 'glu-FpC/bias',
    name: 'glu-FpC/bias_2',
    trainable_: true,
    constraint: null,
    val: Variable {
      kept: false,
      isDisposedInternal: false,
      shape: [Array],
      dtype: 'float32',
      size: 333,
      strides: [],
      dataId: {},
      id: 49,
      rankType: '1',
      trainable: true,
      name: 'glu-FpC/bias_2'
    }
  }

The end result causes issues when attempting to restore those child layers from a checkpoint. I thought it might be possible to fix this with the getWeights() and setWeights() method, but couldn't find an immediate solution:

    getWeights() {
        return this.trainableWeights.map((weights) => weights.read())
    }

    setWeights(weights) {
        this.inProj.kernel.write(weights[0])
        this.inProj.bias.write(weights[1])
        this.outProj.kernel.write(weights[2])
        this.outProj.bias.write(weights[3])
    }

Anyway, I hope that additional context helps.

If you can point me to somewhere I could contribute to the docs, I might be able to write a tutorial or something. I've probably built 100 custom layers at this point.

from tfjs.

Thanks for the suggestion, @lukemovement. Your suggestion does work, however you also need to build each of the child layers:

this.inProj.build(inputShape)
this.outProj.build(inputShape)

I actually discovered this trick previously, and it also has problems. The issue relates to restoring a model from disk:

/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:273
        var _this = _super.call(this, message) || this;
                           ^

ValueError: Duplicate weight name: glu-TLd/kernel
    at new ValueError (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:273:28)
    at Container.loadWeights (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:21823:35)
    at /home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:25792:27
    at step (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:159:27)
    at Object.next (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:108:53)
    at fulfilled (/home/crow/Repos/ode/node_modules/@tensorflow/tfjs-layers/dist/tf-layers.node.js:89:28)

Node.js v18.18.2

When you add child layers to trainableWeights like this, every child layer will inherit the name of its parent (with a suffix):

  LayerVariable {
    dtype: 'float32',
    shape: [ 333 ],
    id: 30,
    originalName: 'glu-FpC/bias',
    name: 'glu-FpC/bias_2',
    trainable_: true,
    constraint: null,
    val: Variable {
      kept: false,
      isDisposedInternal: false,
      shape: [Array],
      dtype: 'float32',
      size: 333,
      strides: [],
      dataId: {},
      id: 49,
      rankType: '1',
      trainable: true,
      name: 'glu-FpC/bias_2'
    }
  }

The end result causes issues when attempting to restore those child layers from a checkpoint. I thought it might be possible to fix this with the getWeights() and setWeights() method, but couldn't find an immediate solution:

    getWeights() {
        return this.trainableWeights.map((weights) => weights.read())
    }

    setWeights(weights) {
        this.inProj.kernel.write(weights[0])
        this.inProj.bias.write(weights[1])
        this.outProj.kernel.write(weights[2])
        this.outProj.bias.write(weights[3])
    }

Anyway, I hope that additional context helps.

If you can point me to somewhere I could contribute to the docs, I might be able to write a tutorial or something. I've probably built 100 custom layers at this point.

I'm unsure as to where this would best be placed. The weights are always in the same order on the layer so I use this. It works as long as you don't reach the maximum string length.

import * as tf from "@tensorflow/tfjs";
import { mkdir, readFile, writeFile } from "fs/promises";
import { resolve } from "path";

export const SaveModel = async ({
  model,
  dir,
}: {
  model: tf.LayersModel;
  dir: string;
}) => {
  for (const layer of model.layers) {
    const name = layer.name;

    if (0 === layer.weights.length) {
      continue;
    }

    const weights = JSON.stringify(
      layer.getWeights().map((weight) => weight.arraySync()),
    );

    await mkdir(dir, { recursive: true });
    await writeFile(resolve(dir, `${name}.json`), weights);
  }

  console.log(`Saved to ${dir}`);
};

export const LoadModel = async ({
  model,
  dir,
}: {
  model: tf.LayersModel;
  dir: string;
}) => {
  for (const layer of model.layers) {
    const name = layer.name;

    if (0 === layer.weights.length) {
      continue;
    }

    try {
      const weights = JSON.parse(
        await readFile(resolve(dir, `${name}.json`), "utf-8"),
      );

      const tensors = weights.map((weight: number[]) => tf.tensor(weight));

      layer.setWeights(tensors);
    } catch (e) {
      console.log(layer.name, (e as Error).message);
    }
  }

  console.warn(`Loaded from ${dir}`);
};

from tfjs.