Custom Layer with kwargs in tfjs

Question

I'm new to tensorflowjs and I'm struggling to implement some custom layers, if someone could point me in the right direction that would be really helpful! For example, I have a layer in InceptionResnetV1 architecture where I'm multiplying the layer by a constant scale (this was originally an unsupported Lambda layer which I'm switching out for a custom layer), but the value of this scale changes per block. This works fine in Keras with an implementation such as below, and using load_model with ScaleLayer in the custom objects

class ScaleLayer(tensorflow.keras.layers.Layer):
    def __init__(self, **kwargs):
      super(ScaleLayer, self).__init__(**kwargs)

    def call(self, inputs, **kwargs):
      return tensorflow.multiply(inputs, kwargs.get('scale'))

    def get_config(self):
        return {}

x = ScaleLayer()(x, scale = tensorflow.constant(scale))

I tried defining this in a similar way in javascript and then registered the class

class ScaleLayer extends tf.layers.Layer {

  constructor(config?: any) {
    super(config || {});
  }

  call(input: tf.Tensor, kwargs: Kwargs) {
    return tf.tidy(() => {
      this.invokeCallHook(input, kwargs);
      const a = input;
      const b = kwargs['scale'];
      return tf.mul(a, b);
    });
  }
  
  static get className() {
    return 'ScaleLayer';
  }
}

tf.serialization.registerClass(ScaleLayer);

However I'm finding that the kwargs are always empty. I tried another similar method where I passed scale as another dimension of the input, then did input[0] * input[1], which again worked fine for the keras model but not in javascript. I feel like I'm missing something key on the way to defining this kind of custom layer with a changing value per block on the javascript end, so if someone would be able to point me in the right direction it would be much appreciated! Thanks.

Answer 1

constructor(config?: any) {
    super(config || {});
  }

The config are passed to the parent constructor. But as indicated by the question, the ScaleLayer layer also needs to keep some config properties

constructor(config?: any) {
    super(config || {});
    // this.propertyOfInterest = config.propertyOfInterest 
    // make sure that config is an object;
    this.scale = config.scale
  }

Then for the computation, the ScaleLayer property propertyOfInterest can be used

call(input: tf.Tensor) {
    return tf.tidy(() => {
      this.invokeCallHook(input, kwargs);
      const a = input;
      return tf.mul(a, this.scale);
    });
  }

Use the layer this way:

const model = tf.sequential();
...
model.add(new ScaleLayer({scale: 1}));
...