transcranial / keras-js

Run Keras models in the browser, with GPU support using WebGL
https://transcranial.github.io/keras-js
MIT License
4.96k stars 503 forks source link

Add a compatibility parameter to the model loader #116

Open linusmartensson opened 6 years ago

linusmartensson commented 6 years ago

So, I had a problem with two models that led me to implement this. In the first one, I had a free-form batch shape, which isn't compatible with keras-js, and in the second one, automatic parameter renaming in newer versions of Keras, which renamed a node from "class_scores_1" to "class_scores", but maintained the original weight name, "class_scores_1/kernel:0", was unable to load.

By implementing this adjustment layer, I can identify such error conditions, and manually add any necessary corrections, without having to wait for such inconsistencies to be resolved, or mess around trying different versions:

const detector = new KerasJS.Model({
    filepath: 'detector.bin',
    gpu:true,
    adjustment: (config)=>{
        console.dir("adjusting...");
        config.config.layers[0].config.batch_input_shape[0] = 1;
        config.config.layers[0].config.batch_input_shape[1] = 360;
        config.config.layers[0].config.batch_input_shape[2] = 640;
        console.dir(config);
        return config;
    }

});
const segmentation = new KerasJS.Model({
    filepath: 'jnet.bin',
    gpu:false,
    adjustment: (config)=>{

        config.config.layers[33].config.name = "class_scores_1";
        config.config.layers[33].name = "class_scores_1";
        config.config.layers[34].inbound_nodes[0][0][0] = "class_scores_1";

        return config;
    }
});

Having this may allow users to work more consistently with the still somewhat experimental API, and I'm thinking it may be useful on a broader scale. For me, it was critical to get the implementation running.