Custom Layer with kwargs in tfjs - tensorflow.js

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}));
...