TensorFlow.js: "Error: Weights are not set" when loading converted model with BatchNormalization Layers

Question

I'm porting to TF.js a model converted from Keras with its pretrained weights. I followed the instructions on the documentation and saved the weights from a local path.

When I'm initializing my model by loading the .json file, with tfjs-node 1.7.1, I get the following error:

(node:39703) UnhandledPromiseRejectionWarning: Error: 152 of 197 weights are not set: conv1_1/3x3_s1/bn/gamma,conv1_1/3x3_s1/bn/beta,conv1_1/3x3_s1/bn/moving_mean,conv1_1/3x3_s1/bn/moving_variance,conv2_a_1x1_reduce/bn/gamma,conv2_a_1x1_reduce/bn/beta,conv2_a_1x1_reduce/bn/moving_mean,conv2_a_1x1_reduce/bn/moving_variance,conv2_a_3x3/bn/gamma,conv2_a_3x3/bn/beta,conv2_a_3x3/bn/moving_mean,conv2_a_3x3/bn/moving_variance,conv2_a_1x1_increase/bn/gamma,conv2_a_1x1_increase/bn/beta,conv2_a_1x1_increase/bn/moving_mean,conv2_a_1x1_increase/bn/moving_variance,conv2_a_1x1_proj/bn/gamma,conv2_a_1x1_proj/bn/beta,conv2_a_1x1_proj/bn/moving_mean,conv2_a_1x1_proj/bn/moving_variance,conv2_b_1x1_reduce/bn/gamma,conv2_b_1x1_reduce/bn/beta,conv2_b_1x1_reduce/bn/moving_mean,conv2_b_1x1_reduce/bn/moving_variance,conv2_b_3x3/bn/gamma,conv2_b_3x3/bn/beta,conv2_b_3x3/bn/moving_mean,conv2_b_3x3/bn/moving_variance,conv2_b_1x1_increase/bn/gamma,conv2_b_1x1_increase/bn/beta,conv2_b_1x1_increase/bn/moving_mean,conv2_b_1x1_increase/bn/moving_variance,conv3_a_1x1_reduce/bn/gamma,conv3_a_1x1_reduce/bn/beta,conv3_a_1x1_reduce/bn/moving_mean,conv3_a_1x1_reduce/bn/moving_variance,conv3_a_3x3/bn/gamma,conv3_a_3x3/bn/beta,conv3_a_3x3/bn/moving_mean,conv3_a_3x3/bn/moving_variance,conv3_a_1x1_increase/bn/gamma,conv3_a_1x1_increase/bn/beta,conv3_a_1x1_increase/bn/moving_mean,conv3_a_1x1_increase/bn/moving_variance,conv3_a_1x1_proj/bn/gamma,conv3_a_1x1_proj/bn/beta,conv3_a_1x1_proj/bn/moving_mean,conv3_a_1x1_proj/bn/moving_variance,conv3_b_1x1_reduce/bn/gamma,conv3_b_1x1_reduce/bn/beta,conv3_b_1x1_reduce/bn/moving_mean,conv3_b_1x1_reduce/bn/moving_variance,conv3_b_3x3/bn/gamma,conv3_b_3x3/bn/beta,conv3_b_3x3/bn/moving_mean,conv3_b_3x3/bn/moving_variance,conv3_b_1x1_increase/bn/gamma,conv3_b_1x1_increase/bn/beta,conv3_b_1x1_increase/bn/moving_mean,conv3_b_1x1_increase/bn/moving_variance,conv3_c_1x1_reduce/bn/gamma,conv3_c_1x1_reduce/bn/beta,conv3_c_1x1_reduce/bn/moving_mean,conv3_c_1x1_reduce/bn/moving_variance,conv3_c_3x3/bn/gamma,conv3_c_3x3/bn/beta,conv3_c_3x3/bn/moving_mean,conv3_c_3x3/bn/moving_variance,conv3_c_1x1_increase/bn/gamma,conv3_c_1x1_increase/bn/beta,conv3_c_1x1_increase/bn/moving_mean,conv3_c_1x1_increase/bn/moving_variance,conv4_a_1x1_reduce/bn/gamma,conv4_a_1x1_reduce/bn/beta,conv4_a_1x1_reduce/bn/moving_mean,conv4_a_1x1_reduce/bn/moving_variance,conv4_a_3x3/bn/gamma,conv4_a_3x3/bn/beta,conv4_a_3x3/bn/moving_mean,conv4_a_3x3/bn/moving_variance,conv4_a_1x1_increase/bn/gamma,conv4_a_1x1_increase/bn/beta,conv4_a_1x1_increase/bn/moving_mean,conv4_a_1x1_increase/bn/moving_variance,conv4_a_1x1_proj/bn/gamma,conv4_a_1x1_proj/bn/beta,conv4_a_1x1_proj/bn/moving_mean,conv4_a_1x1_proj/bn/moving_variance,conv4_b_1x1_reduce/bn/gamma,conv4_b_1x1_reduce/bn/beta,conv4_b_1x1_reduce/bn/moving_mean,conv4_b_1x1_reduce/bn/moving_variance,conv4_b_3x3/bn/gamma,conv4_b_3x3/bn/beta,conv4_b_3x3/bn/moving_mean,conv4_b_3x3/bn/moving_variance,conv4_b_1x1_increase/bn/gamma,conv4_b_1x1_increase/bn/beta,conv4_b_1x1_increase/bn/moving_mean,conv4_b_1x1_increase/bn/moving_variance,conv4_c_1x1_reduce/bn/gamma,conv4_c_1x1_reduce/bn/beta,conv4_c_1x1_reduce/bn/moving_mean,conv4_c_1x1_reduce/bn/moving_variance,conv4_c_3x3/bn/gamma,conv4_c_3x3/bn/beta,conv4_c_3x3/bn/moving_mean,conv4_c_3x3/bn/moving_variance,conv4_c_1x1_increase/bn/gamma,conv4_c_1x1_increase/bn/beta,conv4_c_1x1_increase/bn/moving_mean,conv4_c_1x1_increase/bn/moving_variance,conv5_a_1x1_reduce/bn/gamma,conv5_a_1x1_reduce/bn/beta,conv5_a_1x1_reduce/bn/moving_mean,conv5_a_1x1_reduce/bn/moving_variance,conv5_a_3x3/bn/gamma,conv5_a_3x3/bn/beta,conv5_a_3x3/bn/moving_mean,conv5_a_3x3/bn/moving_variance,conv5_a_1x1_increase/bn/gamma,conv5_a_1x1_increase/bn/beta,conv5_a_1x1_increase/bn/moving_mean,conv5_a_1x1_increase/bn/moving_variance,conv5_a_1x1_proj/bn/gamma,conv5_a_1x1_proj/bn/beta,conv5_a_1x1_proj/bn/moving_mean,conv5_a_1x1_proj/bn/moving_variance,conv5_b_1x1_reduce/bn/gamma,conv5_b_1x1_reduce/bn/beta,conv5_b_1x1_reduce/bn/moving_mean,conv5_b_1x1_reduce/bn/moving_variance,conv5_b_3x3/bn/gamma,conv5_b_3x3/bn/beta,conv5_b_3x3/bn/moving_mean,conv5_b_3x3/bn/moving_variance,conv5_b_1x1_increase/bn/gamma,conv5_b_1x1_increase/bn/beta,conv5_b_1x1_increase/bn/moving_mean,conv5_b_1x1_increase/bn/moving_variance,conv5_c_1x1_reduce/bn/gamma,conv5_c_1x1_reduce/bn/beta,conv5_c_1x1_reduce/bn/moving_mean,conv5_c_1x1_reduce/bn/moving_variance,conv5_c_3x3/bn/gamma,conv5_c_3x3/bn/beta,conv5_c_3x3/bn/moving_mean,conv5_c_3x3/bn/moving_variance,conv5_c_1x1_increase/bn/gamma,conv5_c_1x1_increase/bn/beta,conv5_c_1x1_increase/bn/moving_mean,conv5_c_1x1_increase/bn/moving_variance
    at new ValueError (./AudioReco/node_modules/@tensorflow/tfjs-layers/dist/errors.js:68:28)
    at LayersModel.Container.loadWeights (./AudioReco/node_modules/@tensorflow/tfjs-layers/dist/engine/container.js:569:23)
    at ./AudioReco/node_modules/@tensorflow/tfjs-layers/dist/models.js:303:27
    at step (./AudioReco/node_modules/@tensorflow/tfjs-layers/dist/models.js:54:23)
    at Object.next (./AudioReco/node_modules/@tensorflow/tfjs-layers/dist/models.js:35:53)
    at fulfilled (./AudioReco/node_modules/@tensorflow/tfjs-layers/dist/models.js:26:58)

I understand this has something to do with all the Batch Normalization layers the model uses but how do I properly restore the Gammas, Betas, Moving Mean & Moving Average to each of them?

---

Model taken from the work of Wei Xie from VoxCeleb. The original weights of the model can be found on his Google Drive and the converted ones in the .json format can be found on this link.

Code ran:

model.js

const tf = require('@tensorflow/tfjs-node')

class VGGVox_Model {
    constructor() {
        this.model;
    }

    async init() {
        // Loading the custom layers
        require('./layers/VladPooling');
        require('./layers/Lambda');

        this.model = await tf.loadLayersModel('file://resources/model/model.json', false);
        this.model.summary();
    }
}

(async function main() {
  const myModel = new VGGVox_Model();
  myModel.init();
})();

VladPooling.js

const tf = require('../node_modules/@tensorflow/tfjs-node')

class VladPooling extends tf.layers.Layer {
    constructor(config) {
        super(config);
        this.kCenters = config.kCenters;
        this.gCenters = config.gCenters;
        this.mode = config.mode;
    }


    compute_output_shape(input_shape) {
        return (input_shape[0][0], this.kCenters * input_shape[0][input_shape[0].length - 1])
    }

    build(input_shape) {
        this.cluster = this.addWeight(
            'centers',
            [ this.kCenters + this.gCenters, input_shape[0][input_shape[0].length - 1] ],
            'float32',
            'orthogonal');
        this.built = true;
    }


    call(inputs, kwargs) {
        return tf.tidy(() => {
            console.log('call')
        });
    }


    getConfig() {
        const baseConfig = super.getConfig();
        const config = {
            kCenters: this.kCenters,
            gCenters: this.gCenters,
            mode: this.mode
        };
        Object.assign(config, baseConfig);
        return config;
    }


    static get className() {
        return 'VladPooling';
    }
}
tf.serialization.registerClass(VladPooling);
exports.vlad_pooling = VladPooling;

Lambda.js

const tf = require('../node_modules/@tensorflow/tfjs-node')

class Lambda extends tf.layers.Layer {
    constructor(config) {
        super(config);
        if (config.name === undefined) {
            config.name = ((+new Date) * Math.random()).toString(36); //random name from timestamp in case name hasn't been set
        }
        this.name = config.name;
        this.lambdaFunction = config.function;
    }

    call(input) {
        return tf.tidy(() => {
            let result = null;
            eval(this.lambdaFunction);
            return result;
        });
    }

    computeOutputShape(inputShape) {
        return inputShape;
    }

    getConfig() {
        const config = super.getConfig();
        Object.assign(config, {
            lambdaFunction: this.lambdaFunction
        });
        return config;
    }

    static get className() {
        return 'Lambda';
    }
}
tf.serialization.registerClass(Lambda);

As for the two custom layers (Vlad Pooling and Lambda), I added them as well in the Google Drive folder. Lambda layer's code was taken here. I didn't include the call() of the Vlad Pooling yet since I couldn't check the outputs (model doesn't load). Anyways, they are required because the model loading won't work without them.

Answer 1

Looking at the call method of the VladPoolingLayer in python and js, the js method does not fully implement the VLadPoolingLayer of the python model. call is the method where the actual computation of the layer takes place. Consequently it should return a tensor whose shape matches what is returned by computeOutputShape (typo for VladPoolingLayer).

The lambda layer written was copied from the post here. There is a difference in the way the lambda layer was written compared to a normal lambda layer in python. So this lambda layer cannot be used as it is for the model here. The lambda layer of the VGG-Speaker-Recognition is doing a l2 normalization. Thus the layer can be written as following:

class lambdaLayer extends tf.layers.Layer {
  constructor(config) {
    super(config)
  }

  call(input) {
    return tf.layers.layerNormalization({axis: 1}).apply(input);
  }

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

tf.serialization.registerClass(lambdaLayer);

Answer 2

I had the same issue with YOLOv3 implemented in Keras. I successfully converted the keras model to tfjs and tried to load it using tf.loadLayersModel in a React app but it gave me the exact same error. I thought that it has something to do with the Batch Normalization layers but it was not the case.

The problem was actually with loading the model using loadLayersModel which I don't really know why it is. What I did is that I loaded the model using Keras and saved that model as .pb graph. Then I loaded the .pb using Tensorflow which was working fine. Next step was to convert the model using tensorflowjs_converter wrapper as tfjs_graph_model instead of tfjs_layers_model. Finally I tried to load the tfjs model using tf.loadGraphModel in React app and I got successful making the predictions. The code for converting .h5 model to .pb graph is given below:

import tensorflow as tf
from tensorflow.keras.models import load_model, save_model
from keras import backend as K

def swish_activation(x):
    return (K.sigmoid(x) * x)

keras_model_path = "exported-models/model.h5/"
output_model_path = "exported-models/tf/"

def convert_keras_to_graph():
    model = load_model(keras_model_path, custom_objects={"hard_swish": swish_activation})
    model.summary()
    save_model(model, output_model_path)
    
    # check the model in tensorflow
    model = tf.saved_model.load(output_model_path)


if __name__ == "__main__":
    convert_keras_to_graph()

Now convert the graph in tensorflowjs format using the following cmd command:

tensorflowjs_converter --input_format=tf_saved_model --output_format=tfjs_graph_model --saved_model_tags=serve --signature_name=serving_default exported-models\tf\ exported-models\tfjs --signature_name=serving_default

then in tensorflowjs, just load the model as given below:

const MODEL_URL: "http://localhost:8080/model/model.json",
const model = tf.loadGraphModel(MODEL_URL)
      .then((model) => {
        console.log("model loaded: ", model);
      })
      .catch((error) => {
        console.log("failed to load the model", error);
        throw error;
      });

Note: I've used http-server to place my model files including the weights (.bin) and model.json. This is a good practice to do so.

I hope this will work for you as it has worked for me; upvote this answer in case it has helped you so other can take benefit from it.