Writing custom InstantLayerNormalization in tensorflow js

user_11599950 0
Viewed 74

I am trying to implement a deep learning model in the browser and this requires porting some custom layers, one of them is an instant layer normalization. Below the piece of code that is supposed to work but it's a bit old. I get this error:

Uncaught (in promise) ReferenceError: initializer is not defined at InstantLayerNormalization.build

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs/dist/tf.min.js"> </script>
<script>
class InstantLayerNormalization extends tf.layers.Layer
{
    static className = 'InstantLayerNormalization';
    epsilon = 1e-7 
    gamma;
    beta;
    constructor(config) 
    {
        super(config);
    }
    getConfig() 
    {
        const config = super.getConfig();
        return config;
    }
    
    build(input_shape)
    {
        let shape = tf.tensor(input_shape);
        // initialize gamma
        self.gamma = self.add_weight(shape=shape, 
                                     initializer='ones', 
                                     trainable=true, 
                                     name='gamma')
        // initialize beta
        self.beta = self.add_weight(shape=shape,
                            initializer='zeros',
                            trainable=true,
                            name='beta')
    }        

    call(inputs){
        mean = tf.math.reduce_mean(inputs, axis=[-1], keepdims=True)
        variance = tf.math.reduce_mean(tf.math.square(inputs - mean), axis=[-1], keepdims=True)
        std = tf.math.sqrt(variance + self.epsilon)
        outputs = (inputs - mean) / std
        outputs = outputs * self.gamma
        outputs = outputs + self.beta
        return outputs
    }
    static get className() {
        console.log(className);
       return className;
    }
}

tf.serialization.registerClass(InstantLayerNormalization);
</script>
tensorflow.js tensorflow machine-learning javascript
1 Answers

The methods of the inherited class tf.layers.Layer are not called properly.

  • self in python is this in js
  • add_weight is rather addWeight
  • Here is the signature of the addWeight method. Please notice that in js there is not the format variable=value for function arguments destructuring assignment
// instead of this
self.gamma = self.add_weight(shape=shape, initializer='ones', trainable=true, name='gamma')
// it should rather be
this.gamma = this.addWeight('gamma', shape, undefined, 'ones', undefined, true)
user_50699570
Related