Implementing custom InstanceNormalization layer in tensorflowjs

TensorFlow General Discussion
, ,
1

I am trying to import and load a trained keras cycleGAN model into tensorflowjs and angular typescript.
but i ran into the error that says “ERROR Error: Uncaught (in promise): Error: Unknown layer: InstanceNormalization”

I googled and tried implementing my own custom layer as a .js file:

import * as tf from'@tensorflow/tfjs';

class InstanceNormalization extends tf.layers.Layer {
    epsilon = 1e-7;
    gamma;
    beta;

    constructor() {
        super({});
    }

    build(input_shape) {
        let shape = tf.tensor(input_shape);
        // initialize gamma
        this.gamma = this.addWeight(
            'gamma',
            shape,
            undefined,
            'ones',
            undefined,
            true,
        );

        // initialize beta
        this.beta = this.addWeight(
            'beta',
            shape,
            undefined,
            'zeros',
            undefined,
            true,
        )
    }

    call(inputs, kwargs) {
        this.invokeCallHook(inputs, kwargs);
        let mean = tf.math.reduce_mean(inputs, axis=[-1], keepdims=true);
        let variance = tf.math.reduce_mean(tf.math.square(inputs - mean), axis=[-1], keepdims=true);
        let std = tf.math.sqrt(variance + this.epsilon);
        let outputs = (inputs - mean) / std;
        outputs = outputs * this.gamma;
        outputs = outputs + this.beta;
        return outputs;
    }

    static get className() {
        return 'InstanceNormalization';
    }

}

tf.serialization.registerClass(InstanceNormalization);

export function instancenormalization() {
    return new InstanceNormalization();
}

the above effort did not resolve the issue where the error message still persisted.
these are my references:

  1. javascript - Writing custom InstantLayerNormalization in tensorflow js - Stack Overflow
  2. tfjs-examples/custom-layer/custom_layer.js at master · tensorflow/tfjs-examples · GitHub
3

@pyu or @lina128 May have some thoughts on this one.

4

@G11 Your code looks like having all the plumbings, not sure why it is working for your model. Do you have a codepen example that can reproduce this error? thanks

1 Like